backend.torch_backend#

PyTorch backend — implements AbstractBackend using PyTorch.

The implementation is split across same-package modules by operation category (capabilities, creation, indexing, reductions, passthrough, linalg, interpolation, random, misc); this module composes them into the concrete TorchBackend class. See capabilities.py etc. for the actual method bodies, and config.py for the GradMode/ _Config helpers.

Kramer Harrison, 2025

Classes

TorchBackend()

Backend implementation using PyTorch.

class TorchBackend[source]#

Backend implementation using PyTorch.

Variables:

_lib – The torch module (used by passthrough methods).
_config – Internal configuration (device, precision, grad_mode).

abs(x: Any) → Tensor#

Compute the absolute value of x.

Parameters:: x – Input data.
Returns:: Absolute values.
Return type:: Tensor

all(x: Any) → bool#

Return True if all elements of x are True.

Parameters:: x – Input tensor or bool.
Returns:: Whether all elements are True.
Return type:: bool

allclose(a: Any, b: Any, rtol: float = 1e-05, atol: float = 1e-08) → bool#

Return True if all elements in a and b are close.

Parameters:

a – First input.
b – Second input.
rtol – Relative tolerance.
atol – Absolute tolerance.

Returns:

Whether all elements are close.

Return type:

bool

any(x: Any) → bool#

Return True if any element of x is True.

Parameters:: x – Input tensor or bool.
Returns:: Whether any element is True.
Return type:: bool

arange(*args: Any, **kwargs: Any) → Tensor#

Return evenly spaced values within a given interval.

Parameters:

*args – start, stop, step (positional).
**kwargs – Keyword arguments passed to torch.arange.

Returns:

Evenly spaced values.

Return type:

Tensor

arange_indices(start: Any, stop: Any = None, step: int = 1) → Tensor#

Create a tensor of integer indices.

Parameters:

start – Start index (or stop if stop is None).
stop – Stop index.
step – Step size.

Returns:

Long integer index tensor.

Return type:

Tensor

arccos(x: Any) → Tensor#

Compute the arccosine of x.

Parameters:: x – Input data.
Returns:: Arccosine values.
Return type:: Tensor

arcsin(x: Any) → Tensor#

Compute the arcsine of x.

Parameters:: x – Input data.
Returns:: Arcsine values.
Return type:: Tensor

arctan(x: Any) → Tensor#

Compute the arctangent of x.

Parameters:: x – Input data.
Returns:: Arctangent values.
Return type:: Tensor

arctan2(y: Any, x: Any) → Tensor#

Compute the element-wise arctan2.

Parameters:

y – y-coordinates.
x – x-coordinates.

Returns:

Arctangent values (angle in radians).

Return type:

Tensor

argmin(x: Any, axis: int | None = None) → Tensor#

Return indices of the minimum values along a dimension.

Parameters:

x – Input tensor.
axis – Dimension along which to find the minimum.

Returns:

Index tensor.

Return type:

Tensor

argwhere(x: Any) → Tensor#

Return indices of non-zero elements.

Parameters:: x – Input tensor.
Returns:: Index tensor of shape (N, ndim).
Return type:: Tensor

array(x: ArrayLike) → Tensor#

Create a tensor with current device, precision, and grad settings.

Parameters:: x – Input data.
Returns:: Backend tensor.
Return type:: Tensor

as_array_1d(data: Any) → Tensor#

Force conversion to a 1-D tensor.

Parameters:: data – Scalar, list, tuple, or tensor.
Returns:: 1-D tensor.
Return type:: Tensor
Raises:: ValueError – If data type is not supported.

asarray(x: Any, **kwargs: Any) → Tensor#

Convert x to a tensor without copying if possible.

Parameters:

x – Input data.
**kwargs – Keyword arguments forwarded to torch.as_tensor (e.g. dtype). NumPy dtypes are automatically converted to the equivalent torch dtype.

Returns:

Backend tensor.

Return type:

Tensor

atleast_1d(x: Any) → Tensor#

Convert x to a tensor with at least one dimension.

Parameters:: x – Input data.
Returns:: At least 1-D tensor.
Return type:: Tensor

atleast_2d(x: Any) → Tensor#

Convert x to a tensor with at least two dimensions.

Parameters:: x – Input data.
Returns:: At least 2-D tensor.
Return type:: Tensor

property autograd: Any#: Return the torch.autograd submodule.

batched_chain_matmul3(a: Tensor, b: Tensor, c: Tensor) → Tensor#

Compute a @ b @ c with promoted dtype.

Parameters:

a – First matrix.
b – Second matrix.
c – Third matrix.

Returns:

Result of a @ b @ c.

Return type:

Tensor

broadcast_to(x: Tensor, shape: Sequence[int]) → Tensor#

Broadcast x to the given shape.

Parameters:

x – Input tensor.
shape – Target shape.

Returns:

Broadcast tensor.

Return type:

Tensor

cast(x: Any) → Tensor#

Cast x to the current floating-point precision.

Parameters:: x – Input data.
Returns:: Cast tensor.
Return type:: Tensor

ceil(x: Any) → Tensor#

Round up to the nearest integer.

Parameters:: x – Input data.
Returns:: Ceiling values.
Return type:: Tensor

clip(x: Any, a_min: Any, a_max: Any) → Tensor#

Clip values in x to [a_min, a_max].

Parameters:

x – Input tensor.
a_min – Minimum value.
a_max – Maximum value.

Returns:

Clipped tensor.

Return type:

Tensor

column_stack(*args: Any, **kwargs: Any) → Any#

concatenate(arrays: Sequence[Any], axis: int = 0) → Tensor#

Join tensors along an existing axis.

Parameters:

arrays – Sequence of tensors to concatenate.
axis – Axis along which to concatenate.

Returns:

Concatenated tensor.

Return type:

Tensor

conj(x: Any) → Tensor#

Compute the complex conjugate of x.

Parameters:: x – Input data.
Returns:: Complex conjugate.
Return type:: Tensor

copy(x: Tensor) → Tensor#

Return a copy of x.

Parameters:: x – Input tensor.
Returns:: Cloned tensor.
Return type:: Tensor

copy_to(source: Tensor, destination: Tensor) → None#

In-place copy from source to destination tensor.

Safely handles tensors that require gradients.

Parameters:

source – Source tensor.
destination – Destination tensor (modified in place).

copysign(x: Any, y: Any) → Tensor#

Return x with the sign of y (element-wise).

Parameters:

x – Magnitude array.
y – Sign array.

Returns:

x with sign from y.

Return type:

Tensor

cos(x: Any) → Tensor#

Compute the cosine of x (element-wise).

Parameters:: x – Input data.
Returns:: Cosine values.
Return type:: Tensor

cosh(x: Any) → Tensor#

Compute the hyperbolic cosine of x.

Parameters:: x – Input data.
Returns:: Hyperbolic cosine values.
Return type:: Tensor

cross(a: Tensor, b: Tensor, axisa: int = -1, axisb: int = -1, axisc: int = -1, axis: int | None = None) → Tensor#

Return the cross product of two vectors.

Parameters:

a – First vector tensor.
b – Second vector tensor.
axisa – Axis of a defining the vector(s).
axisb – Axis of b defining the vector(s).
axisc – Axis of c containing the cross product.
axis – If set, applies to axisa, axisb, and axisc.

Returns:

Cross product.

Return type:

Tensor

default_rng(seed: int | None = None) → TorchGenerator#

Return a PyTorch random number generator.

Parameters:: seed – Optional seed.
Returns:: PyTorch Generator.
Return type:: Generator

deg2rad(x: Any) → Tensor#

Convert angles from degrees to radians.

Parameters:: x – Angle in degrees.
Returns:: Angle in radians.
Return type:: Tensor

degrees(x: Any) → Tensor#

Convert angles from radians to degrees.

Parameters:: x – Angle in radians.
Returns:: Angle in degrees.
Return type:: Tensor

diff(x: Any, n: int = 1, axis: int = -1, **kwargs: Any) → Tensor#

Calculate the n-th discrete difference along a dimension.

Parameters:

x – Input tensor.
n – Number of times to apply the difference.
axis – Dimension along which to compute differences.
**kwargs – Additional keyword arguments forwarded to torch.diff (e.g. prepend, append).

Returns:

Differences tensor.

Return type:

Tensor

dot(*args: Any, **kwargs: Any) → Any#

einsum(*args: Any, **kwargs: Any) → Any#

empty(shape: Sequence[int]) → Tensor#

Return an uninitialized tensor of given shape.

Parameters:: shape – Shape of the output tensor.
Returns:: Uninitialized tensor.
Return type:: Tensor

empty_like(x: Any) → Tensor#

Return an uninitialized tensor with the same shape as x.

Parameters:: x – Reference tensor.
Returns:: Uninitialized tensor.
Return type:: Tensor

erfinv(x: Any) → Tensor#

Inverse error function.

Parameters:: x – Input tensor.
Returns:: Inverse error function of x.
Return type:: Tensor

errstate(**kwargs: Any) → Generator[None, None, None]#

No-op context manager (torch has no equivalent of np.errstate).

Parameters:: **kwargs – Ignored.
Yields:: None

exp(x: Any) → Tensor#

Compute the exponential of x.

Parameters:: x – Input data.
Returns:: Exponential values.
Return type:: Tensor

expand_dims(x: Tensor, axis: int) → Tensor#

Insert a new axis into x.

Parameters:

x – Input tensor.
axis – Position of the new axis.

Returns:

Tensor with new axis.

Return type:

Tensor

eye(n: int) → Tensor#

Return a 2D identity matrix.

Parameters:: n – Size of the matrix.
Returns:: Identity matrix.
Return type:: Tensor

factorial(n: Any) → Tensor#

Compute the factorial of n using the log-gamma function.

Parameters:: n – Non-negative integer or tensor.
Returns:: Factorial values.
Return type:: Tensor

property fft: Any#: Expose the FFT submodule of the underlying library.

fftconvolve(in1: Tensor, in2: Tensor, mode: Literal['full', 'valid', 'same'] = 'full') → Tensor#

FFT-based convolution using PyTorch.

Parameters:

in1 – First input tensor (N-D).
in2 – Second input tensor (N-D).
mode – Convolution mode ('full', 'valid', 'same').

Returns:

Convolved tensor.

Return type:

Tensor

Raises:

ValueError – If inputs have different dimensionality or mode is unknown.

finfo(*args: Any, **kwargs: Any) → Any#

flip(x: Tensor) → Tensor#

Reverse the order of elements along axis 0.

Parameters:: x – Input tensor.
Returns:: Flipped tensor.
Return type:: Tensor

floor(x: Any) → Tensor#

Round down to the nearest integer.

Parameters:: x – Input data.
Returns:: Floor values.
Return type:: Tensor

fmax(a: Any, b: Any) → Tensor#

Element-wise maximum, ignoring NaNs.

Parameters:

a – First input tensor.
b – Second input tensor.

Returns:

Element-wise maximum ignoring NaN.

Return type:

Tensor

full(shape: Sequence[int], fill_value: Any, dtype: Any = None) → Tensor#

Return a constant-filled tensor of given shape.

Parameters:

shape – Shape of the output tensor.
fill_value – Fill value.
dtype – Optional dtype override.

Returns:

Filled tensor.

Return type:

Tensor

full_like(x: Any, fill_value: Any) → Tensor#

Return a full tensor with the same shape as x.

Parameters:

x – Reference tensor.
fill_value – Fill value.

Returns:

Filled tensor. When fill_value is itself a Tensor,: the result stays attached to its autograd graph (via a broadcast multiply) instead of being detached through .item() – a scalar requires_grad leaf broadcast across a multi-element x must still carry a gradient back to that leaf.

Return type:

Tensor

get_bilinear_weights(coords: Tensor, bin_edges: Sequence[Tensor]) → tuple[Tensor, Tensor]#

Compute differentiable bilinear interpolation weights.

Parameters:

coords – Ray coordinates tensor of shape (N, 2).
bin_edges – Sequence of two edge tensors [x_edges, y_edges].

Returns:

(all_indices, all_weights).

Return type:

tuple[Tensor, Tensor]

get_complex_precision() → torch.dtype#

Return the complex dtype matching the current float precision.

Returns:: torch.complex64 or torch.complex128.
Return type:: torch.dtype
Raises:: ValueError – If the current precision is unsupported.

get_device() → str#: Return the current compute device.

get_precision() → int#: Return the current precision as an integer (32 or 64).

property grad_mode: GradMode#: Return the GradMode controller.

grid_sample(input: Tensor, grid: Tensor, mode: str = 'bilinear', padding_mode: str = 'zeros', align_corners: bool = False) → Tensor#

Sample input using torch.nn.functional.grid_sample.

Parameters:

input – Input tensor of shape (N, C, H_in, W_in).
grid – Grid tensor of shape (N, H_out, W_out, 2).
mode – Interpolation mode.
padding_mode – Padding mode.
align_corners – Whether to align corners.

Returns:

Output tensor of shape (N, C, H_out, W_out).

Return type:

Tensor

histogram(x: Any, bins: Any = 10) → tuple[Tensor, Tensor]#

Compute a histogram of x.

Parameters:

x – Input tensor.
bins – Number of bins or bin edge tensor.

Returns:

Bin counts and bin edges.

Return type:

tuple[Tensor, Tensor]

histogram2d(x: Tensor, y: Tensor, bins: Any, weights: Tensor | None = None) → tuple[Tensor, Tensor, Tensor]#

Compute a 2-D histogram.

Parameters:

x – x-coordinates of the sample points.
y – y-coordinates of the sample points.
bins – List or tuple of two edge tensors.
weights – Optional weights for each sample.

Returns:

Histogram, x edges, y edges.

Return type:

tuple[Tensor, Tensor, Tensor]

Raises:

ValueError – If bins is not a list/tuple of two edge tensors.

hypot(x: Any, y: Any) → Tensor#

Compute the hypotenuse given legs x and y.

Parameters:

x – First leg.
y – Second leg.

Returns:

Hypotenuse values.

Return type:

Tensor

imag(x: Any) → Tensor#

Return the imaginary part of x.

Parameters:: x – Input data.
Returns:: Imaginary part.
Return type:: Tensor

index_add(target: Tensor, dim: int, index: Tensor, source: Tensor) → Tensor#

Scatter-add: result[index[i]] += source[i] along dim.

Parameters:

target – Accumulation buffer.
dim – Dimension along which to scatter.
index – 1-D integer indices tensor.
source – Values to add (same dtype as target).

Returns:

New tensor with values scatter-added (non-in-place).

Return type:

Tensor

interp(x: Any, xp: Any, fp: Any) → Tensor#

1-D linear interpolation.

Parameters:

x – x-coordinates of the interpolated values.
xp – x-coordinates of the data points.
fp – y-coordinates of the data points.

Returns:

Interpolated values.

Return type:

Tensor

is_array_like(x: Any) → bool#

Return True if x is a tensor, ndarray, list, or tuple.

Parameters:: x – Object to check.
Returns:: True if x is array-like.
Return type:: bool

isclose(a: Any, b: Any, rtol: float = 1e-05, atol: float = 1e-08) → Tensor#

Return a boolean tensor where elements are close.

Parameters:

a – First input.
b – Second input.
rtol – Relative tolerance.
atol – Absolute tolerance.

Returns:

Boolean tensor.

Return type:

Tensor

isfinite(x: Any) → Any#

Check if input is finite, accepting scalars and tensors.

Parameters:: x – Input (scalar, ndarray, or Tensor).
Returns:: Whether x is finite.
Return type:: bool or Tensor

isinf(x: Any) → Any#

Check if input is infinity, accepting scalars and tensors.

Parameters:: x – Input (scalar, ndarray, or Tensor).
Returns:: Whether x is infinite.
Return type:: bool or Tensor

isnan(x: Any) → Any#

Check if input is NaN, accepting scalars and tensors.

Parameters:: x – Input (scalar, ndarray, or Tensor).
Returns:: Whether x is NaN.
Return type:: bool or Tensor

isscalar(x: Any) → bool#

Return True if x is a 0-dimensional tensor.

Parameters:: x – Input.
Returns:: Whether x is a scalar tensor.
Return type:: bool

property linalg: Any#: Expose the linear-algebra submodule of the underlying library.

linspace(start: float, stop: float, num: int = 50) → Tensor#

Return evenly spaced numbers over an interval.

Parameters:

start – Start of the interval.
stop – End of the interval.
num – Number of samples.

Returns:

Evenly spaced samples.

Return type:

Tensor

load(filename: str) → Tensor#

Load a NumPy file and convert to a tensor.

Parameters:: filename – Path to a .npy file.
Returns:: Loaded tensor.
Return type:: Tensor

log(x: Any) → Tensor#

Compute the natural logarithm of x.

Parameters:: x – Input data.
Returns:: Natural log values.
Return type:: Tensor

log10(x: Any) → Tensor#

Compute the base-10 logarithm of x.

Parameters:: x – Input data.
Returns:: log10 values.
Return type:: Tensor

log2(x: Any) → Tensor#

Compute the base-2 logarithm of x.

Parameters:: x – Input data.
Returns:: log2 values.
Return type:: Tensor

logical_and(*args: Any, **kwargs: Any) → Any#

logical_not(*args: Any, **kwargs: Any) → Any#

logical_or(*args: Any, **kwargs: Any) → Any#

lstsq(a: Tensor, b: Tensor) → Tensor#

Compute the least-squares solution to a @ x = b.

Parameters:

a – Left-hand side matrix (M, N).
b – Right-hand side matrix (M,) or (M, K).

Returns:

Least-squares solution.

Return type:

Tensor

matmul(a: Tensor, b: Tensor) → Tensor#

Matrix product of two tensors with promoted dtype.

Parameters:

a – First matrix.
b – Second matrix.

Returns:

Matrix product.

Return type:

Tensor

matrix_vector_multiply_and_squeeze(p: Tensor, E: Tensor) → Tensor#

Multiply p @ E[…, newaxis] and squeeze trailing dimension.

Parameters:

p – Matrix tensor.
E – Vector tensor.

Returns:

Result with trailing dimension squeezed.

Return type:

Tensor

max(x: Any) → Any#

Return the maximum value of x.

Parameters:: x – Input tensor or array.
Returns:: Maximum value as a Python scalar.
Return type:: float

maximum(a: Any, b: Any) → Tensor#

Element-wise maximum of a and b.

Parameters:

a – First input tensor.
b – Second input tensor.

Returns:

Element-wise maximum.

Return type:

Tensor

mean(x: Any, axis: int | None = None, keepdims: bool = False) → Tensor#

Compute the arithmetic mean, ignoring NaNs.

Parameters:

x – Input tensor.
axis – Dimension along which to compute the mean.
keepdims – Whether to keep reduced dimensions.

Returns:

Mean.

Return type:

Tensor

meshgrid(*arrays: Tensor) → tuple[Tensor, ...]#

Return coordinate matrices from coordinate vectors (xy indexing).

Parameters:: *arrays – 1-D tensors representing grid coordinates.
Returns:: Coordinate matrices.
Return type:: tuple[Tensor, …]

min(x: Any) → Any#

Return the minimum value of x.

Parameters:: x – Input tensor or array.
Returns:: Minimum value as a Python scalar.
Return type:: float

minimum(a: Any, b: Any) → Tensor#

Element-wise minimum of a and b.

Parameters:

a – First input tensor.
b – Second input tensor.

Returns:

Element-wise minimum.

Return type:

Tensor

mult_p_E(p: Tensor, E: Tensor) → Tensor#

Complex matrix-vector multiply for polarized fields.

Parameters:

p – Jones matrix tensor.
E – Electric field tensor.

Returns:

Complex matrix-vector product.

Return type:

Tensor

property name: str#: Return the backend name.

nanmax(x: Tensor, axis: int | None = None, keepdim: bool = False) → Tensor#

Return the maximum value, ignoring NaNs.

Parameters:

x – Input tensor.
axis – Dimension along which to compute the maximum.
keepdim – Whether to keep reduced dimensions.

Returns:

Maximum value ignoring NaN.

Return type:

Tensor

nanmean(*args: Any, **kwargs: Any) → Any#

nansum(*args: Any, **kwargs: Any) → Any#

nearest_nd_interpolator(points: Tensor, values: Tensor, Hx: Tensor, Hy: Tensor) → Tensor#

Nearest-neighbour interpolation on an N-D dataset.

Parameters:

points – Known sample points of shape (N, D).
values – Values at the sample points.
Hx – Query x coordinates.
Hy – Query y coordinates.

Returns:

Interpolated values.

Return type:

Tensor

Raises:

ValueError – If Hx or Hy is None.

property nn: Any#: Return the torch.nn submodule.

ones(shape: Sequence[int], dtype: Any = None) → Tensor#

Return a ones tensor of given shape.

Parameters:

shape – Shape of the output tensor.
dtype – Optional dtype override.

Returns:

Ones tensor.

Return type:

Tensor

ones_like(x: Any) → Tensor#

Return a ones tensor with the same shape as x.

Parameters:: x – Reference tensor.
Returns:: Ones tensor.
Return type:: Tensor

outer(*args: Any, **kwargs: Any) → Any#

pad(tensor: Tensor, pad_width: Any, mode: str = 'constant', constant_values: float | None = 0) → Tensor#

Pad a tensor.

Parameters:

tensor – Input tensor.
pad_width – Padding per axis.
mode – Padding mode.
constant_values – Value used for constant padding.

Returns:

Padded tensor.

Return type:

Tensor

path_contains_points(vertices: Tensor, points: Tensor) → Tensor#

Return a boolean mask of points inside the polygon.

Uses a vectorized ray-crossing algorithm.

Parameters:

vertices – Polygon vertices as (N, 2) tensor (closed implicitly).
points – Query points as (M, 2) tensor.

Returns:

Boolean mask of shape (M,).

Return type:

Tensor

polyfit(x: Tensor, y: Tensor, degree: int) → Tensor#

Least-squares polynomial fit.

Parameters:

x – x-coordinates of the sample points.
y – y-coordinates of the sample points.
degree – Degree of the polynomial.

Returns:

Polynomial coefficients, highest power first.

Return type:

Tensor

polyval(coeffs: Any, x: Any) → Any#

Evaluate a polynomial at specific values.

Parameters:

coeffs – Polynomial coefficients, highest power first.
x – Values at which to evaluate.

Returns:

Evaluated polynomial.

Return type:

Tensor or float

power(x: Any, y: Any) → Tensor#

Return x raised to the power y.

Parameters:

x – Base tensor.
y – Exponent tensor.

Returns:

x ** y.

Return type:

Tensor

rad2deg(x: Any) → Tensor#

Convert angles from radians to degrees.

Parameters:: x – Angle in radians.
Returns:: Angle in degrees.
Return type:: Tensor

radians(x: Any) → Tensor#

Convert angles from degrees to radians.

Parameters:: x – Angle in degrees.
Returns:: Angle in radians.
Return type:: Tensor

rand(*size: int) → Tensor#

Random values from a uniform distribution on [0, 1).

Parameters:: *size – Shape of the output tensor.
Returns:: Random values.
Return type:: Tensor

property random: Any#: Expose the random submodule of the underlying library.

random_normal(loc: float = 0.0, scale: float = 1.0, size: Any = None, generator: TorchGenerator | None = None) → Tensor#

Random samples from a Gaussian distribution.

Parameters:

loc – Mean of the distribution.
scale – Standard deviation.
size – Output shape.
generator – Optional torch Generator.

Returns:

Normal random samples.

Return type:

Tensor

random_uniform(low: float = 0.0, high: float = 1.0, size: Any = None, generator: TorchGenerator | None = None) → Tensor#

Uniform random samples in [low, high).

Parameters:

low – Lower boundary.
high – Upper boundary.
size – Output shape.
generator – Optional torch Generator.

Returns:

Uniform random samples.

Return type:

Tensor

ravel(x: Any) → Tensor#

Return a contiguous flattened tensor.

Parameters:: x – Input tensor, array, or scalar.
Returns:: Flattened tensor.
Return type:: Tensor

real(x: Any) → Tensor#

Return the real part of x.

Parameters:: x – Input data.
Returns:: Real part.
Return type:: Tensor

repeat(x: Tensor, repeats: int) → Tensor#

Repeat elements of x.

Parameters:

x – Input tensor.
repeats – Number of repetitions.

Returns:

Repeated tensor.

Return type:

Tensor

reshape(x: Tensor, shape: Sequence[int]) → Tensor#

Return x with a new shape.

Parameters:

x – Input tensor.
shape – New shape.

Returns:

Reshaped tensor.

Return type:

Tensor

roll(x: Tensor, shift: Any, axis: Any = ()) → Tensor#

Roll tensor elements along the given axis.

Parameters:

x – Input tensor.
shift – Number of places to shift.
axis – Axis or axes along which to roll.

Returns:

Rolled tensor.

Return type:

Tensor

round(*args: Any, **kwargs: Any) → Any#

searchsorted(*args: Any, **kwargs: Any) → Any#

set_device(device: str) → None#

Set the compute device.

Parameters:: device – 'cpu' or 'cuda'.

set_precision(precision: Literal['float32', 'float64']) → None#

Set the floating-point precision.

Parameters:: precision – 'float32' or 'float64'.

shape(tensor: Tensor) → tuple[int, ...]#

Return the shape of a tensor.

Parameters:: tensor – Input tensor.
Returns:: Shape of the tensor.
Return type:: tuple[int, …]

sign(x: Any) → Tensor#

Compute the sign of x.

Parameters:: x – Input data.
Returns:: Sign values (-1, 0, or 1).
Return type:: Tensor

sin(x: Any) → Tensor#

Compute the sine of x (element-wise).

Parameters:: x – Input data.
Returns:: Sine values.
Return type:: Tensor

sinh(x: Any) → Tensor#

Compute the hyperbolic sine of x.

Parameters:: x – Input data.
Returns:: Hyperbolic sine values.
Return type:: Tensor

size(x: Any) → int#

Return the total number of elements in x.

Parameters:: x – Input tensor, array, or scalar.
Returns:: Number of elements.
Return type:: int

sobol_sampler(dim: int, num_samples: int, scramble: bool = True, seed: int | None = None) → Tensor#

Generate quasi-random samples using Sobol sequences.

Parameters:

dim – Dimension of the samples.
num_samples – Number of samples to generate.
scramble – Whether to scramble the sequence.
seed – Random seed for scrambling.

Returns:

Samples of shape (num_samples_pow2, dim).

Return type:

Tensor

sort(x: Any, axis: int = -1) → Tensor#

Return a sorted tensor along the given dimension.

Parameters:

x – Input tensor.
axis – Dimension along which to sort.

Returns:

Sorted tensor (values only, not the indices).

Return type:

Tensor

sqrt(x: Any) → Tensor#

Compute the square root of x.

Parameters:: x – Input data.
Returns:: Square root values.
Return type:: Tensor

stack(xs: Sequence[Any], axis: int = 0) → Tensor#

Join a sequence of tensors along a new axis.

Parameters:

xs – Sequence of tensors.
axis – Axis along which to stack.

Returns:

Stacked tensor.

Return type:

Tensor

std(x: Any, axis: int | None = None) → Tensor#

Compute the standard deviation along an axis.

Parameters:

x – Input tensor.
axis – Dimension along which to compute std.

Returns:

Standard deviation.

Return type:

Tensor

sum(x: Any, axis: int | None = None) → Tensor#

Sum tensor elements over a given axis.

Parameters:

x – Input tensor.
axis – Dimension along which to sum.

Returns:

Sum.

Return type:

Tensor

property supports_gpu: bool#: Return True if CUDA is available.

property supports_gradients: bool#: Return True — PyTorch supports automatic differentiation.

tan(x: Any) → Tensor#

Compute the tangent of x (element-wise).

Parameters:: x – Input data.
Returns:: Tangent values.
Return type:: Tensor

tanh(x: Any) → Tensor#

Compute the hyperbolic tangent of x.

Parameters:: x – Input data.
Returns:: Hyperbolic tangent values.
Return type:: Tensor

tensor(data: Any, **kwargs: Any) → Tensor#

Create a tensor from data with full kwargs support.

Parameters:

data – Input data (scalar, list, numpy array, etc.).
**kwargs – Forwarded to torch.tensor (e.g. requires_grad, dtype, device).

Returns:

New tensor.

Return type:

Tensor

tile(x: Tensor, dims: Any) → Tensor#

Construct a tensor by tiling x.

Parameters:

x – Input tensor.
dims – Number of repetitions per dimension.

Returns:

Tiled tensor.

Return type:

Tensor

to_complex(x: Tensor) → Tensor#

Cast x to complex128.

Parameters:: x – Input tensor.
Returns:: Complex128 tensor.
Return type:: Tensor

to_tensor(data: ArrayLike, device: str | torch.device | None = None) → Tensor#

Convert data to a PyTorch tensor with the backend’s precision.

Parameters:

data – The data to convert.
device – Optional device override.

Returns:

Converted tensor.

Return type:

Tensor

transpose(x: Tensor, axes: Sequence[int] | None = None) → Tensor#

Permute the dimensions of x.

Parameters:

x – Input tensor.
axes – The dimensions to permute. If None, reverses all dimensions.

Returns:

Transposed tensor.

Return type:

Tensor

unsqueeze_last(x: Tensor) → Tensor#

Add a trailing dimension to x.

Parameters:: x – Input tensor.
Returns:: Tensor with extra trailing dimension.
Return type:: Tensor

vectorize(pyfunc: Callable[..., Any]) → Callable[..., Any]#

Vectorize a scalar Python function over tensor inputs.

Parameters:: pyfunc – The scalar function to vectorize.
Returns:: Vectorized function.
Return type:: Callable

vstack(*args: Any, **kwargs: Any) → Any#

where(condition: Any, x: Any, y: Any) → Any#

Return elements from x or y depending on condition.

Parameters:

condition – Boolean tensor or bool.
x – Values where condition is True.
y – Values where condition is False.

Returns:

Output tensor.

Return type:

Tensor

zeros(shape: Sequence[int], dtype: Any = None) → Tensor#

Return a zero tensor of given shape.

Parameters:

shape – Shape of the output tensor.
dtype – Optional dtype override.

Returns:

Zero tensor.

Return type:

Tensor

zeros_like(x: Any) → Tensor#

Return a zero tensor with the same shape as x.

Parameters:: x – Reference tensor.
Returns:: Zero tensor.
Return type:: Tensor