backend.torch_backend

PyTorch backend — implements AbstractBackend using PyTorch.

Kramer Harrison, 2025

Classes

GradMode()

Control global gradient computation for the torch backend.

TorchBackend()

Backend implementation using PyTorch.

class GradMode[source]

Control global gradient computation for the torch backend.

disable() None[source]

Disable gradient computation.

enable() None[source]

Enable gradient computation.

temporary_enable() Generator[None, None, None][source]

Context manager that temporarily enables gradient computation.

class TorchBackend[source]

Backend implementation using PyTorch.

_lib

The torch module (used by passthrough methods).

_config

Internal configuration (device, precision, grad_mode).

abs(x: Any) Tensor[source]

Compute the absolute value of x.

Parameters:

x – Input data.

Returns:

Absolute values.

Return type:

Tensor

all(x: Any) bool[source]

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[source]

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[source]

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[source]

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[source]

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[source]

Compute the arccosine of x.

Parameters:

x – Input data.

Returns:

Arccosine values.

Return type:

Tensor

arcsin(x: Any) Tensor[source]

Compute the arcsine of x.

Parameters:

x – Input data.

Returns:

Arcsine values.

Return type:

Tensor

arctan(x: Any) Tensor[source]

Compute the arctangent of x.

Parameters:

x – Input data.

Returns:

Arctangent values.

Return type:

Tensor

arctan2(y: Any, x: Any) Tensor[source]

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[source]

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[source]

Return indices of non-zero elements.

Parameters:

x – Input tensor.

Returns:

Index tensor of shape (N, ndim).

Return type:

Tensor

array(x: ArrayLike) Tensor[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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[source]

Broadcast x to the given shape.

Parameters:
  • x – Input tensor.

  • shape – Target shape.

Returns:

Broadcast tensor.

Return type:

Tensor

cast(x: Any) Tensor[source]

Cast x to the current floating-point precision.

Parameters:

x – Input data.

Returns:

Cast tensor.

Return type:

Tensor

ceil(x: Any) Tensor[source]

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[source]

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[source]

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[source]

Compute the complex conjugate of x.

Parameters:

x – Input data.

Returns:

Complex conjugate.

Return type:

Tensor

copy(x: Tensor) Tensor[source]

Return a copy of x.

Parameters:

x – Input tensor.

Returns:

Cloned tensor.

Return type:

Tensor

copy_to(source: Tensor, destination: Tensor) None[source]

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[source]

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[source]

Compute the cosine of x (element-wise).

Parameters:

x – Input data.

Returns:

Cosine values.

Return type:

Tensor

cosh(x: Any) Tensor[source]

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[source]

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[source]

Return a PyTorch random number generator.

Parameters:

seed – Optional seed.

Returns:

PyTorch Generator.

Return type:

Generator

deg2rad(x: Any) Tensor[source]

Convert angles from degrees to radians.

Parameters:

x – Angle in degrees.

Returns:

Angle in radians.

Return type:

Tensor

degrees(x: Any) Tensor[source]

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[source]

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[source]

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[source]

Return an uninitialized tensor with the same shape as x.

Parameters:

x – Reference tensor.

Returns:

Uninitialized tensor.

Return type:

Tensor

erfinv(x: Any) Tensor[source]

Inverse error function.

Parameters:

x – Input tensor.

Returns:

Inverse error function of x.

Return type:

Tensor

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

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

Parameters:

**kwargs – Ignored.

Yields:

None

exp(x: Any) Tensor[source]

Compute the exponential of x.

Parameters:

x – Input data.

Returns:

Exponential values.

Return type:

Tensor

expand_dims(x: Tensor, axis: int) Tensor[source]

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[source]

Return a 2D identity matrix.

Parameters:

n – Size of the matrix.

Returns:

Identity matrix.

Return type:

Tensor

factorial(n: Any) Tensor[source]

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[source]

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[source]

Reverse the order of elements along axis 0.

Parameters:

x – Input tensor.

Returns:

Flipped tensor.

Return type:

Tensor

floor(x: Any) Tensor[source]

Round down to the nearest integer.

Parameters:

x – Input data.

Returns:

Floor values.

Return type:

Tensor

fmax(a: Any, b: Any) Tensor[source]

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[source]

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[source]

Return a full tensor with the same shape as x.

Parameters:
  • x – Reference tensor.

  • fill_value – Fill value.

Returns:

Filled tensor.

Return type:

Tensor

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

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[source]

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[source]

Return the current compute device.

get_precision() int[source]

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[source]

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][source]

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][source]

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[source]

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[source]

Return the imaginary part of x.

Parameters:

x – Input data.

Returns:

Imaginary part.

Return type:

Tensor

interp(x: Any, xp: Any, fp: Any) Tensor[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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[source]

Compute the natural logarithm of x.

Parameters:

x – Input data.

Returns:

Natural log values.

Return type:

Tensor

log10(x: Any) Tensor[source]

Compute the base-10 logarithm of x.

Parameters:

x – Input data.

Returns:

log10 values.

Return type:

Tensor

log2(x: Any) Tensor[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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, ...][source]

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[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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[source]

Pad a tensor.

Parameters:
  • tensor – Input tensor.

  • pad_width – Padding per axis as ((pt, pb), (pl, pr)).

  • mode – Only 'constant' is supported.

  • constant_values – Value used for constant padding.

Returns:

Padded tensor.

Return type:

Tensor

Raises:

NotImplementedError – If mode is not 'constant'.

path_contains_points(vertices: Tensor, points: Tensor) Tensor[source]

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[source]

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[source]

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[source]

Return x raised to the power y.

Parameters:
  • x – Base tensor.

  • y – Exponent tensor.

Returns:

x ** y.

Return type:

Tensor

rad2deg(x: Any) Tensor[source]

Convert angles from radians to degrees.

Parameters:

x – Angle in radians.

Returns:

Angle in degrees.

Return type:

Tensor

radians(x: Any) Tensor[source]

Convert angles from degrees to radians.

Parameters:

x – Angle in degrees.

Returns:

Angle in radians.

Return type:

Tensor

rand(*size: int) Tensor[source]

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[source]

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[source]

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: Tensor) Tensor[source]

Return a contiguous flattened tensor.

Parameters:

x – Input tensor.

Returns:

Flattened tensor.

Return type:

Tensor

real(x: Any) Tensor[source]

Return the real part of x.

Parameters:

x – Input data.

Returns:

Real part.

Return type:

Tensor

repeat(x: Tensor, repeats: int) Tensor[source]

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[source]

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[source]

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[source]

Set the compute device.

Parameters:

device'cpu' or 'cuda'.

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

Set the floating-point precision.

Parameters:

precision'float32' or 'float64'.

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

Return the shape of a tensor.

Parameters:

tensor – Input tensor.

Returns:

Shape of the tensor.

Return type:

tuple[int, …]

sign(x: Any) Tensor[source]

Compute the sign of x.

Parameters:

x – Input data.

Returns:

Sign values (-1, 0, or 1).

Return type:

Tensor

sin(x: Any) Tensor[source]

Compute the sine of x (element-wise).

Parameters:

x – Input data.

Returns:

Sine values.

Return type:

Tensor

sinh(x: Any) Tensor[source]

Compute the hyperbolic sine of x.

Parameters:

x – Input data.

Returns:

Hyperbolic sine values.

Return type:

Tensor

size(x: Tensor) int[source]

Return the total number of elements in x.

Parameters:

x – Input tensor.

Returns:

Number of elements.

Return type:

int

sobol_sampler(dim: int, num_samples: int, scramble: bool = True, seed: int | None = None) Tensor[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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[source]

Compute the tangent of x (element-wise).

Parameters:

x – Input data.

Returns:

Tangent values.

Return type:

Tensor

tanh(x: Any) Tensor[source]

Compute the hyperbolic tangent of x.

Parameters:

x – Input data.

Returns:

Hyperbolic tangent values.

Return type:

Tensor

tensor(data: Any, **kwargs: Any) Tensor[source]

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[source]

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[source]

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[source]

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[source]

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[source]

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][source]

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[source]

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[source]

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[source]

Return a zero tensor with the same shape as x.

Parameters:

x – Reference tensor.

Returns:

Zero tensor.

Return type:

Tensor