latest (dev)
Copy
Latest development documentation · Updated 2026-10-15
Source code for tensorplay._composite_funcs
"""
Complements ``functional.py`` (codegen) and ``_shape_funcs.py``. Every
public name here is expressed over existing differentiable primitives
(mm/cat/narrow/gather/index_select/pow/sort/searchsorted/bincount/...),
so autograd flows through composition wherever the underlying ops are
``test/test_composite_funcs.py``.
Deliberate narrows (documented per function):
- conjugation-bit queries report physical layout (tp has no conj bit);
- ``cartesian_prod`` accepts 1-D inputs.
"""
import itertools
import math
import numpy as np
import tensorplay
from tensorplay._C import DType
from tensorplay.autograd.function import Function
from tensorplay.graph import capture_call as _capture_call
__all__ = [
"absolute", "arccos", "arccosh", "arcsin", "arcsinh", "arctan",
"arctanh", "arctan2",
"acos_", "asin_", "atan_", "acosh_", "asinh_", "atanh_",
"concat", "concatenate", "ger", "adjoint",
"detach", "diagflat", "numel", "scalar_tensor",
"chain_matmul", "matrix_power", "kron", "vander",
"tril_indices", "triu_indices", "cartesian_prod", "combinations",
"trapezoid", "trapz", "cumulative_trapezoid",
"quantile", "nanquantile", "histc", "histogram",
"isin", "unique_consecutive", "repeat_interleave", "kaiser_window",
"lstm_cell", "rnn_relu_cell", "rnn_tanh_cell",
"put", "resolve_conj", "resolve_neg", "is_conj", "is_neg",
"can_cast", "promote_types", "result_type", "is_nonzero",
"is_same_size", "get_device",
"alias_copy", "t_copy", "permute_copy", "transpose_copy",
"squeeze_copy", "unsqueeze_copy", "select_copy", "slice_copy",
"narrow_copy", "diagonal_copy", "unbind_copy", "split_copy",
"view_copy", "unfold_copy", "expand_copy", "reshape_as",
"unsafe_chunk", "unsafe_split",
"rms_norm", "cosine_similarity",
"max_pool1d", "avg_pool1d", "adaptive_avg_pool1d",
"adaptive_max_pool1d",
"dropout", "dropout_", "alpha_dropout", "feature_dropout",
"feature_dropout_", "feature_alpha_dropout", "rrelu", "rrelu_",
"bilinear", "ctc_loss", "embedding_bag", "conv_tbc",
"max_pool3d", "max_pool1d_with_indices", "native_channel_shuffle",
]
# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------
def _norm_dim(dim, ndim):
d = int(dim)
if d < 0:
d += ndim
if d < 0 or d >= ndim:
raise IndexError(
f"Dimension out of range (expected to be in range of "
f"[{-ndim}, {ndim - 1}], but got {dim})"
)
return d
def _as_tensor(x):
if isinstance(x, tensorplay.Tensor):
return x
return tensorplay.as_tensor(x)
def _band(a, b):
return tensorplay.logical_and(a, b)
def _slice_getitem(t, dim, start, end, step=1):
sl = [slice(None)] * t.dim()
sl[dim] = slice(int(start), int(end), int(step))
return t[tuple(sl)]
def _make_arc_inplace(base, deriv):
class _ArcInplaceFn(Function):
@staticmethod
def forward(ctx, self):
ctx.pre = self.clone()
res = getattr(self, base)()
with tensorplay.no_grad():
self.copy_(res)
return self
@staticmethod
def backward(ctx, grad_output):
return (grad_output * deriv(ctx.pre),)
def op(self):
return _ArcInplaceFn.apply(self)
op.__name__ = base + "_"
return op
def _d_acos(x):
return tensorplay.full_like(x, 1.0).sub_(x.mul(x)).rsqrt().neg()
def _d_asin(x):
return tensorplay.full_like(x, 1.0).sub_(x.mul(x)).rsqrt()
def _d_atan(x):
return x.mul(x).add(1.0).reciprocal()
def _d_acosh(x):
return x.mul(x).sub(1.0).rsqrt()
def _d_asinh(x):
return x.mul(x).add(1.0).rsqrt()
def _d_atanh(x):
return x.neg().mul(x).add(1.0).reciprocal()
acos_ = _make_arc_inplace("acos", _d_acos)
asin_ = _make_arc_inplace("asin", _d_asin)
atan_ = _make_arc_inplace("atan", _d_atan)
acosh_ = _make_arc_inplace("acosh", _d_acosh)
asinh_ = _make_arc_inplace("asinh", _d_asinh)
atanh_ = _make_arc_inplace("atanh", _d_atanh)
# ---------------------------------------------------------------------------
# aliases / semantic wrappers
# ---------------------------------------------------------------------------
def absolute(input):
"""Alias of :func:`abs`."""
return input.abs()
def arccos(input):
return input.acos()
def arcsin(input):
return input.asin()
def arctan(input):
return input.atan()
def arccosh(input):
return input.acosh()
def arcsinh(input):
return input.asinh()
def arctanh(input):
return input.atanh()
def arctan2(input, other):
"""atan2(y, x) = 2*atan(y / (hypot + x)); native atan2 CPU kernel is
currently unregistered in this tree, so compose. Edge narrows vs
"""
y = input if isinstance(input, tensorplay.Tensor) else _as_tensor(input)
x = other if isinstance(other, tensorplay.Tensor) else _as_tensor(other)
xd = x.to(DType.float64)
yd = y.to(DType.float64)
r = xd.mul(xd).add(yd.mul(yd)).sqrt()
den = r.add(xd)
den_safe = tensorplay.where(
den.eq(0), tensorplay.full_like(den, 1e-300), den)
base = yd.div(den_safe).atan().mul(2.0)
neg_half_pi = -math.pi / 2.0
left = tensorplay.where(
yd.lt(0),
tensorplay.full_like(base, neg_half_pi),
tensorplay.full_like(base, math.pi / 2.0),
)
out = tensorplay.where(_band(xd.lt(0), yd.eq(0)), left, base)
od = _float_out_dtype(input.dtype)
return out.to(od)
def concat(tensors, dim=0, *, out=None):
result = tensorplay.cat(tensors=list(tensors), dim=dim)
if out is not None:
out.copy_(result)
return out
return result
def concatenate(tensors, dim=0, *, out=None):
return concat(tensors, dim=dim, out=out)
def ger(input, vec2):
return input.outer(vec2)
def adjoint(input):
out = input.transpose(-2, -1)
if input.dtype in (DType.complex64, DType.complex128,
DType.bcomplex32, DType.complex32):
out = tensorplay.conj(out)
return out
def detach(input):
return input.detach()
def diagflat(input, offset=0):
return input.reshape([-1]).diag(offset)
def numel(obj):
return obj.numel()
def scalar_tensor(s, *, dtype=None, device=None):
t = tensorplay.tensor([s], dtype=dtype, device=device)
return t.reshape([])
# ---------------------------------------------------------------------------
# linear algebra / structured
# ---------------------------------------------------------------------------
def chain_matmul(*matrices):
if len(matrices) == 0:
raise RuntimeError("chain_matmul(): expected at least one matrix")
out = matrices[0]
for m in matrices[1:]:
out = out.mm(m)
return out
def matrix_power(input, n):
n = int(n)
if input.dim() != 2:
raise RuntimeError("matrix_power(): expected a 2-D square matrix")
if input.size(0) != input.size(1):
raise RuntimeError("matrix_power(): expected a square matrix")
result = tensorplay.eye(input.size(0), dtype=input.dtype,
device=input.device)
base = tensorplay.linalg.inv(input) if n < 0 else input
e = -n if n < 0 else n
while e > 0:
if e & 1:
result = result.mm(base)
e >>= 1
if e:
base = base.mm(base)
return result
def kron(input, other):
a = input if input.dim() >= 1 else input.reshape([1])
b = other if other.dim() >= 1 else other.reshape([1])
if a.dim() != b.dim():
raise RuntimeError(
f"kron(): number of dimensions must match, got {a.dim()} "
f"and {b.dim()}"
)
ndim = a.dim()
ashape = list(a.size())
bshape = list(b.size())
pa = []
pb = []
merged = []
for s_a, s_b in zip(ashape, bshape):
pa.extend([s_a, 1])
pb.extend([1, s_b])
merged.append(s_a * s_b)
out = a.reshape(pa).mul(b.reshape(pb))
return out.reshape(merged)
def vander(x, N=None, increasing=False):
if x.dim() != 1:
raise RuntimeError("vander(): expected a 1-D input tensor")
n = x.numel() if N is None else int(N)
if n < 0:
raise ValueError("vander(): N must be non-negative")
exps = tensorplay.arange(n, device=x.device)
if not increasing:
exps = tensorplay.flip(exps, dims=[0])
return x.unsqueeze(1).pow(exps.unsqueeze(0))
def tril_indices(row, col, offset=0, *, dtype=DType.int64, device=None,
pin_memory=False):
row, col, offset = int(row), int(col), int(offset)
if row < 0 or col < 0:
raise ValueError("tril_indices(): row and col must be non-negative")
rr = tensorplay.arange(row, device=device).unsqueeze(1).repeat([1, col])
cc = tensorplay.arange(col, device=device).unsqueeze(0).repeat([row, 1])
mask = cc.le(rr + offset)
r_flat = rr.masked_select(mask)
c_flat = cc.masked_select(mask)
out = tensorplay.stack([r_flat, c_flat], dim=0)
if dtype != DType.int64:
out = out.to(dtype)
if pin_memory:
out = out.pin_memory()
return out
def triu_indices(row, col, offset=0, *, dtype=DType.int64, device=None,
pin_memory=False):
row, col, offset = int(row), int(col), int(offset)
if row < 0 or col < 0:
raise ValueError("triu_indices(): row and col must be non-negative")
rr = tensorplay.arange(row, device=device).unsqueeze(1).repeat([1, col])
cc = tensorplay.arange(col, device=device).unsqueeze(0).repeat([row, 1])
mask = cc.ge(rr + offset)
r_flat = rr.masked_select(mask)
c_flat = cc.masked_select(mask)
out = tensorplay.stack([r_flat, c_flat], dim=0)
if dtype != DType.int64:
out = out.to(dtype)
if pin_memory:
out = out.pin_memory()
return out
def cartesian_prod(*tensors):
if len(tensors) == 0:
raise RuntimeError("cartesian_prod(): expected at least one tensor")
ts = [t if t.dim() >= 1 else t.reshape([1]) for t in tensors]
if any(t.dim() != 1 for t in ts):
raise NotImplementedError(
"cartesian_prod(): only 1-D inputs are supported"
)
grids = tensorplay.meshgrid(list(ts), indexing="ij")
cols = [g.reshape([-1, 1]).to(ts[0].dtype) for g in grids]
return tensorplay.cat(cols, dim=1)
def combinations(input, r=2, with_replacement=False):
n = input.size(0)
r = int(r)
if r < 0:
raise ValueError("combinations(): r must be non-negative")
rr = r
gen = (itertools.combinations_with_replacement(range(n), r)
if with_replacement else itertools.combinations(range(n), r))
idx_list = list(gen)
tail = list(input.size())[1:]
if rr == 0:
return tensorplay.zeros([0] + tail, dtype=input.dtype,
device=input.device)
if not idx_list:
return tensorplay.zeros([0] + tail, dtype=input.dtype,
device=input.device)
idx = tensorplay.tensor(idx_list, dtype=DType.int64, device=input.device)
picked = input.index_select(0, idx.reshape([-1]))
return picked.reshape([idx.size(0), rr] + tail)
# ---------------------------------------------------------------------------
# statistics / numerics
# ---------------------------------------------------------------------------
_INT_DTYPES = (DType.uint8, DType.int8, DType.int16, DType.int32,
DType.int64)
def _float_out_dtype(dt):
return DType.float64 if dt == DType.float64 else DType.float32
def trapezoid(y, x=None, *, dx=None, dim=-1):
d = _norm_dim(dim, y.dim())
n = y.size(d)
if n == 0:
if x is not None:
xs = _as_tensor(x)
if xs.dim() == 1 and xs.numel() != 0:
raise RuntimeError(
"trapezoid(): there must be one x value for every "
"sample in y along dim"
)
return y.narrow(d, 0, 0).sum(d)
if x is None:
step = dx if dx is not None else 1.0
avg = y.narrow(d, 0, n - 1).add(y.narrow(d, 1, n - 1)).mul(step / 2.0)
else:
xs = _as_tensor(x)
if xs.dim() == 1:
if xs.numel() != y.size(d):
raise RuntimeError(
"trapezoid(): there must be one x value for every "
"sample in y along dim"
)
view = [1] * y.dim()
view[d] = xs.numel()
xs = xs.reshape(view)
n = y.size(d)
dxs = xs.narrow(d, 1, n - 1).sub(xs.narrow(d, 0, n - 1))
avg = y.narrow(d, 0, n - 1).add(y.narrow(d, 1, n - 1)) \
.mul(0.5).mul(dxs)
return avg.sum(d)
def trapz(y, x=None, *, dx=None, dim=-1):
return trapezoid(y, x, dx=dx, dim=dim)
def cumulative_trapezoid(y, x=None, *, dx=None, dim=-1):
d = _norm_dim(dim, y.dim())
n = y.size(d)
if n == 0:
if x is not None:
xs = _as_tensor(x)
if xs.dim() == 1 and xs.numel() != 0:
raise RuntimeError(
"cumulative_trapezoid(): there must be one x value for "
"every sample in y along dim"
)
return y.narrow(d, 0, 0)
if x is None:
step = dx if dx is not None else 1.0
avg = y.narrow(d, 0, n - 1).add(y.narrow(d, 1, n - 1)).mul(step / 2.0)
else:
xs = _as_tensor(x)
if xs.dim() == 1:
if xs.numel() != y.size(d):
raise RuntimeError(
"cumulative_trapezoid(): there must be one x value for "
"every sample in y along dim"
)
view = [1] * y.dim()
view[d] = xs.numel()
xs = xs.reshape(view)
n = y.size(d)
dxs = xs.narrow(d, 1, n - 1).sub(xs.narrow(d, 0, n - 1))
avg = y.narrow(d, 0, n - 1).add(y.narrow(d, 1, n - 1)) \
.mul(0.5).mul(dxs)
return avg.cumsum(d)
[docs]
def quantile(input, q, dim=None, keepdim=False, *, interpolation="linear"):
"""Compute quantiles through the dispatcher composite implementation.
All five interpolation modes, float32/float64, NaN semantics, and output
shapes are handled here. A Python-number q is
"""
# A Python number becomes a tensor of the input's type; anything else
# (a tensor, or a value standing for one under a trace) goes to the
# binding, which rejects what is not a tensor.
qs = q
if isinstance(q, (int, float)) and not isinstance(q, tensorplay.Tensor):
qs = tensorplay.tensor(q, dtype=input.dtype, device=input.device)
return tensorplay._C.quantile(input, qs, dim=dim, keepdim=keepdim,
interpolation=interpolation)
[docs]
def nanquantile(input, q, dim=None, keepdim=False, *, interpolation="linear"):
# A Python number becomes a tensor of the input's type; anything else
# (a tensor, or a value standing for one under a trace) goes to the
# binding, which rejects what is not a tensor.
qs = q
if isinstance(q, (int, float)) and not isinstance(q, tensorplay.Tensor):
qs = tensorplay.tensor(q, dtype=input.dtype, device=input.device)
return tensorplay._C.nanquantile(input, qs, dim=dim, keepdim=keepdim,
interpolation=interpolation)
def histc(input, bins=100, min=0, max=0):
bins = int(bins)
if bins <= 0:
raise ValueError("histc(): bins must be a positive integer")
v = input.reshape([-1]).to(DType.float64)
lo, hi = float(min), float(max)
if lo == 0 and hi == 0:
if v.numel() == 0:
return tensorplay.zeros(
[bins], dtype=input.dtype, device=input.device)
lo = v.min().item()
hi = v.max().item()
if lo > hi:
raise RuntimeError("histc(): upper bound must be larger than "
"lower bound")
if lo == hi or v.numel() == 0:
return tensorplay.zeros([bins], dtype=input.dtype, device=input.device)
width = (hi - lo) / bins
in_range = _band(v.ge(lo), v.le(hi))
idx = ((v - lo) / width).floor().to(DType.int64) \
.clamp(0, bins - 1).masked_select(in_range)
counts = tensorplay.bincount(idx, minlength=bins).narrow(0, 0, bins)
return counts.to(input.dtype)
[docs]
def histogram(input, bins=100, range=None, *, weight=None, density=False):
"""
composite): hist/bin_edges in the input dtype, linspace edges, aminmax
outer-edge inference with the numpy empty-range expansion, weight in the
input dtype and density normalization."""
if isinstance(bins, tensorplay.Tensor):
return tensorplay._C.histogram(input, bins, weight=weight,
density=density)
return tensorplay._C.histogram(input, bins, range=range, weight=weight,
density=density)
def _weighted_bincount(idx, weight, nbins, density, widths):
if weight is not None:
w = weight.reshape([-1]).to(DType.float64)
acc = tensorplay.zeros(
[nbins], dtype=DType.float64, device=idx.device)
acc = acc.index_put([idx], w, accumulate=True)
else:
acc = tensorplay.bincount(idx, minlength=nbins).narrow(0, 0, nbins) \
.to(DType.float64)
if density:
total = acc.sum().item()
if total > 0:
acc = acc.div(widths.mul(total))
return acc
def isin(elements, test_elements, *, assume_unique=False, invert=False):
el = elements.reshape([-1])
te = test_elements.reshape([-1])
if te.numel() == 0:
out = tensorplay.zeros(
el.numel(), dtype=DType.bool, device=elements.device)
else:
ts, _ = tensorplay.sort(te)
pos = tensorplay.searchsorted(ts, el, right=True)
cand = pos.sub(1).clamp(0, ts.numel() - 1)
found = ts.index_select(0, cand).eq(el)
out = _band(found, pos.ge(1))
out = out.reshape(elements.size())
if invert:
return out.logical_not()
return out
def unique_consecutive(input, return_inverse=False, return_counts=False,
dim=None):
if dim is not None:
dim = _norm_dim(dim, input.dim())
moved = tensorplay.movedim(input, dim, 0)
n = moved.size(0)
if n == 0:
output = tensorplay.empty(
list(input.shape), dtype=input.dtype, device=input.device)
empty_index = tensorplay.empty(
[0], dtype=DType.int64, device=input.device)
if return_inverse and return_counts:
return output, empty_index, empty_index
if return_inverse:
return output, empty_index
if return_counts:
return output, empty_index
return output
flat = moved.reshape([n, -1])
if n == 1:
change = tensorplay.ones(
[1], dtype=DType.bool, device=input.device)
else:
adjacent = flat.narrow(0, 1, n - 1).ne(
flat.narrow(0, 0, n - 1))
change = tensorplay.cat(
[tensorplay.ones([1], dtype=DType.bool, device=input.device),
tensorplay.any(adjacent, dim=1)],
dim=0,
)
kept = tensorplay.nonzero(change).reshape([-1])
num_runs = kept.numel()
output = flat.index_select(0, kept).reshape(
[num_runs] + list(moved.shape[1:]))
output = tensorplay.movedim(output, 0, dim)
inverse = change.to(DType.int64).cumsum(0).sub(1)
counts = tensorplay.zeros(
[num_runs], dtype=DType.int64, device=input.device)
counts = counts.scatter_add(
0, inverse,
tensorplay.ones([n], dtype=DType.int64, device=input.device))
if return_inverse and return_counts:
return output, inverse, counts
if return_inverse:
return output, inverse
if return_counts:
return output, counts
return output
v = input.reshape([-1])
n = v.numel()
if n == 0:
outs = [tensorplay.zeros([0], dtype=input.dtype, device=input.device)]
if return_inverse:
outs.append(tensorplay.zeros([0], dtype=DType.int64,
device=input.device))
if return_counts:
outs.append(tensorplay.zeros([0], dtype=DType.int64,
device=input.device))
return outs[0] if len(outs) == 1 else tuple(outs)
ne = v.narrow(0, 1, n - 1).ne(v.narrow(0, 0, n - 1))
change = tensorplay.cat(
[tensorplay.ones([1], dtype=DType.bool, device=input.device), ne])
gid = change.to(DType.int64).cumsum(0).sub(1)
values = v.masked_select(change)
ngid = int(gid.narrow(0, n - 1, 1).item()) + 1
outs = [values]
if return_inverse:
outs.append(gid)
if return_counts:
counts = tensorplay.zeros(
[ngid], dtype=DType.int64, device=input.device) \
.scatter_add(
0, gid,
tensorplay.ones([n], dtype=DType.int64, device=input.device))
outs.append(counts)
return outs[0] if len(outs) == 1 else tuple(outs)
def repeat_interleave(input, repeats, dim=None, *, output_size=None):
src = input
if dim is None:
src = input.reshape([1, -1])
axis = 1
else:
axis = _norm_dim(dim, input.dim())
# Scalar repeats keep the repetition count on the host, so the result
# shape is known without reading device data; an expanded view plus one
# materializing copy replaces the index-map construction.
if isinstance(repeats, (int, bool)) and not isinstance(
repeats, tensorplay.Tensor):
rep_n = int(repeats)
if rep_n < 0:
raise ValueError(
"repeat_interleave(): repeats must be non-negative")
n = src.size(axis)
total = n * rep_n
if output_size is not None and int(output_size) != total:
raise RuntimeError(
"repeat_interleave(): output_size does not match the repeated "
f"length ({int(output_size)} != {total})"
)
if total == 0:
output_shape = list(src.shape)
output_shape[axis] = 0
picked = tensorplay.empty(
output_shape, dtype=input.dtype, device=input.device)
return picked.reshape([-1]) if dim is None else picked
view_shape = list(src.shape)
view_shape.insert(axis + 1, rep_n)
out_shape = list(src.shape)
out_shape[axis] = total
picked = src.unsqueeze(axis + 1).expand(view_shape) \
.reshape(out_shape)
if dim is None:
return picked.reshape([-1])
return picked
if isinstance(repeats, tensorplay.Tensor) or not isinstance(
repeats, (int, float, bool)):
rep = _as_tensor(repeats).to(DType.int64).to(device=input.device)
rep = rep.reshape([-1])
if rep.numel() == 1 and src.size(axis) != 1:
rep = tensorplay.full([src.size(axis)], int(rep[0].item()),
dtype=DType.int64, device=input.device)
else:
rep = tensorplay.full([src.size(axis)], int(repeats),
dtype=DType.int64, device=input.device)
if rep.numel() != src.size(axis):
raise RuntimeError(
"repeat_interleave(): repeats must have the same length as the "
"selected dimension"
)
if bool((rep < 0).any()):
raise ValueError("repeat_interleave(): repeats must be non-negative")
ends = rep.cumsum(0)
total = 0 if rep.numel() == 0 else int(ends[-1].item())
if output_size is not None and int(output_size) != total:
raise RuntimeError(
"repeat_interleave(): output_size does not match the repeated "
f"length ({int(output_size)} != {total})"
)
if total == 0:
output_shape = list(src.shape)
output_shape[axis] = 0
picked = tensorplay.empty(
output_shape, dtype=input.dtype, device=input.device)
return picked.reshape([-1]) if dim is None else picked
pos = tensorplay.arange(total, dtype=DType.int64, device=input.device)
seg = tensorplay.searchsorted(ends.contiguous(), pos, right=True) \
.clamp(0, src.size(axis) - 1)
moved = src.transpose(0, axis) if axis != 0 else src
picked = moved.index_select(0, seg)
if axis != 0:
picked = picked.transpose(0, axis)
if dim is None:
return picked.reshape([-1])
return picked
def kaiser_window(window_length, periodic=True, beta=12.0, *,
dtype=None, layout=None, device=None, pin_memory=False,
requires_grad=False):
n = int(window_length)
if n < 0:
raise ValueError("kaiser_window(): window_length must be non-negative")
if layout not in (None, tensorplay.strided):
raise ValueError("kaiser_window(): only strided layout is supported")
if dtype is None:
dtype = tensorplay.get_default_dtype()
if n <= 1:
w = tensorplay.ones([n], dtype=dtype, device=device)
if pin_memory:
w = w.pin_memory()
return w.requires_grad_(True) if requires_grad else w
denom = n if periodic else n - 1
if denom > 0:
t = tensorplay.arange(n, dtype=DType.float64, device=device)
alpha = denom / 2.0
x = (t - alpha) / alpha
arg = (1.0 - x.pow(2)).clamp(min=0.0).sqrt().mul(float(beta))
def _i0_series(z):
# I0(z) = sum_k ((z^2)/4)^k / (k!)^2, converged in float64
out = tensorplay.ones_like(z)
term = tensorplay.ones_like(z)
for k in range(1, 40):
term = term.mul(z.mul(z).div(4.0)).div(k * k)
out = out.add(term)
return out
w = _i0_series(arg).div(_i0_series(
tensorplay.full([], float(beta), dtype=DType.float64,
device=device)))
else:
w = tensorplay.ones([n], dtype=DType.float64, device=device)
w = w.to(dtype)
if pin_memory:
w = w.pin_memory()
return w.requires_grad_(True) if requires_grad else w
# ---------------------------------------------------------------------------
# rnn cells
# ---------------------------------------------------------------------------
def lstm_cell(input, hx, cx, w_ih, w_hh, b_ih=None, b_hh=None):
gates = input.mm(w_ih.t()).add(hx.mm(w_hh.t()))
if b_ih is not None:
gates = gates.add(b_ih)
if b_hh is not None:
gates = gates.add(b_hh)
w = gates.size(1) // 4
gi = gates.narrow(1, 0, w)
gf = gates.narrow(1, w, w)
go = gates.narrow(1, 2 * w, w)
gg = gates.narrow(1, 3 * w, w)
i = gi.sigmoid()
f = gf.sigmoid()
g = gg.tanh()
o = go.sigmoid()
cy = f.mul(cx).add(i.mul(g))
hy = o.mul(cy.tanh())
return hy, cy
def rnn_relu_cell(input, hx, w_ih, w_hh, b_ih=None, b_hh=None):
h = input.mm(w_ih.t()).add(hx.mm(w_hh.t()))
if b_ih is not None:
h = h.add(b_ih)
if b_hh is not None:
h = h.add(b_hh)
return h.relu()
def rnn_tanh_cell(input, hx, w_ih, w_hh, b_ih=None, b_hh=None):
h = input.mm(w_ih.t()).add(hx.mm(w_hh.t()))
if b_ih is not None:
h = h.add(b_ih)
if b_hh is not None:
h = h.add(b_hh)
return h.tanh()
# ---------------------------------------------------------------------------
# misc tensor utilities
# ---------------------------------------------------------------------------
def put(input, index, source):
flat = input.clone().reshape([-1])
idx = index.to(DType.int64).reshape([-1])
src = source.to(input.dtype).reshape([-1])
if src.numel() == 0 and idx.numel() > 0:
raise RuntimeError("put(): source must not be empty when indices "
"are given")
if src.numel() < idx.numel():
reps = (idx.numel() + src.numel() - 1) // src.numel()
src = src.repeat([reps])
if src.numel() != idx.numel():
src = src.narrow(0, 0, idx.numel())
flat = flat.index_put([idx], src)
return flat.reshape(list(input.size()))
def resolve_conj(input):
return input
def resolve_neg(input):
return input
def is_conj(input):
return False
def is_neg(input):
return False
_PROMOTE_TABLE = {
"bool": {'bool': 'bool', 'uint8': 'uint8', 'int8': 'int8', 'int16': 'int16', 'int32': 'int32', 'int64': 'int64', 'float16': 'float16', 'bfloat16': 'bfloat16', 'float32': 'float32', 'float64': 'float64', 'complex64': 'complex64', 'complex128': 'complex128'},
"uint8": {'bool': 'uint8', 'uint8': 'uint8', 'int8': 'int16', 'int16': 'int16', 'int32': 'int32', 'int64': 'int64', 'float16': 'float16', 'bfloat16': 'bfloat16', 'float32': 'float32', 'float64': 'float64', 'complex64': 'complex64', 'complex128': 'complex128'},
"int8": {'bool': 'int8', 'uint8': 'int16', 'int8': 'int8', 'int16': 'int16', 'int32': 'int32', 'int64': 'int64', 'float16': 'float16', 'bfloat16': 'bfloat16', 'float32': 'float32', 'float64': 'float64', 'complex64': 'complex64', 'complex128': 'complex128'},
"int16": {'bool': 'int16', 'uint8': 'int16', 'int8': 'int16', 'int16': 'int16', 'int32': 'int32', 'int64': 'int64', 'float16': 'float16', 'bfloat16': 'bfloat16', 'float32': 'float32', 'float64': 'float64', 'complex64': 'complex64', 'complex128': 'complex128'},
"int32": {'bool': 'int32', 'uint8': 'int32', 'int8': 'int32', 'int16': 'int32', 'int32': 'int32', 'int64': 'int64', 'float16': 'float16', 'bfloat16': 'bfloat16', 'float32': 'float32', 'float64': 'float64', 'complex64': 'complex64', 'complex128': 'complex128'},
"int64": {'bool': 'int64', 'uint8': 'int64', 'int8': 'int64', 'int16': 'int64', 'int32': 'int64', 'int64': 'int64', 'float16': 'float16', 'bfloat16': 'bfloat16', 'float32': 'float32', 'float64': 'float64', 'complex64': 'complex64', 'complex128': 'complex128'},
"float16": {'bool': 'float16', 'uint8': 'float16', 'int8': 'float16', 'int16': 'float16', 'int32': 'float16', 'int64': 'float16', 'float16': 'float16', 'bfloat16': 'float32', 'float32': 'float32', 'float64': 'float64', 'complex64': 'complex64', 'complex128': 'complex128'},
"bfloat16": {'bool': 'bfloat16', 'uint8': 'bfloat16', 'int8': 'bfloat16', 'int16': 'bfloat16', 'int32': 'bfloat16', 'int64': 'bfloat16', 'float16': 'float32', 'bfloat16': 'bfloat16', 'float32': 'float32', 'float64': 'float64', 'complex64': 'complex64', 'complex128': 'complex128'},
"float32": {'bool': 'float32', 'uint8': 'float32', 'int8': 'float32', 'int16': 'float32', 'int32': 'float32', 'int64': 'float32', 'float16': 'float32', 'bfloat16': 'float32', 'float32': 'float32', 'float64': 'float64', 'complex64': 'complex64', 'complex128': 'complex128'},
"float64": {'bool': 'float64', 'uint8': 'float64', 'int8': 'float64', 'int16': 'float64', 'int32': 'float64', 'int64': 'float64', 'float16': 'float64', 'bfloat16': 'float64', 'float32': 'float64', 'float64': 'float64', 'complex64': 'complex128', 'complex128': 'complex128'},
"complex64": {'bool': 'complex64', 'uint8': 'complex64', 'int8': 'complex64', 'int16': 'complex64', 'int32': 'complex64', 'int64': 'complex64', 'float16': 'complex64', 'bfloat16': 'complex64', 'float32': 'complex64', 'float64': 'complex128', 'complex64': 'complex64', 'complex128': 'complex128'},
"complex128": {'bool': 'complex128', 'uint8': 'complex128', 'int8': 'complex128', 'int16': 'complex128', 'int32': 'complex128', 'int64': 'complex128', 'float16': 'complex128', 'bfloat16': 'complex128', 'float32': 'complex128', 'float64': 'complex128', 'complex64': 'complex128', 'complex128': 'complex128'},
}
_RANK = {
DType.bool: 0,
DType.uint8: 1, DType.int8: 2, DType.int16: 3, DType.int32: 4,
DType.int64: 5,
DType.float16: 6, DType.bfloat16: 6, DType.float32: 7, DType.float64: 8,
DType.complex64: 9, DType.complex128: 10,
}
_CATEGORY = {
DType.bool: 0,
DType.uint8: 1, DType.int8: 1, DType.int16: 1, DType.int32: 1,
DType.int64: 1,
DType.float16: 2, DType.bfloat16: 2, DType.float32: 2, DType.float64: 2,
DType.complex64: 3, DType.complex128: 3,
}
def _dt(dt):
return dt if isinstance(dt, DType) else DType(dt)
def promote_types(type1, type2):
a, b = _dt(type1), _dt(type2)
an, bn = a.name, b.name
if an in _PROMOTE_TABLE and bn in _PROMOTE_TABLE:
return getattr(DType, _PROMOTE_TABLE[an][bn])
if _CATEGORY.get(a) != _CATEGORY.get(b):
hi = a if _CATEGORY.get(a, 0) > _CATEGORY.get(b, 0) else b
cat = _CATEGORY[hi]
for name in ("float32", "complex64"):
if getattr(DType, name) is not None and \
_CATEGORY.get(getattr(DType, name)) == cat:
pass
return hi
return a if _RANK[a] >= _RANK[b] else b
def can_cast(from_, to):
a, b = _dt(from_), _dt(to)
an, bn = a.name, b.name
if an in _PROMOTE_TABLE and bn in _PROMOTE_TABLE:
return _PROMOTE_TABLE[an][bn] == bn
ca, cb = _CATEGORY[a], _CATEGORY[b]
if ca != cb:
return ca < cb
return _RANK[a] <= _RANK[b]
def result_type(*args):
"""The type an elementwise operation on these operands produces.
Operands rank by category: tensors with dimensions over 0-dim tensors over
plain numbers, and within a category by promotion. A number takes part
as a wrapped number -- its kind counts, never a width: a whole number
beside a 32-bit tensor stays 32-bit, a real number turns a whole tensor
into the default real type, and a complex number keeps a real tensor's
width in the complex result.
"""
dim_result = zero_result = wrapped_result = None
def promote(a, b):
return b if a is None else (a if b is None else promote_types(a, b))
for item in args:
if isinstance(item, tensorplay.Tensor):
if item.dim() > 0:
dim_result = promote(dim_result, item.dtype)
else:
zero_result = promote(zero_result, item.dtype)
continue
if isinstance(item, bool):
dt = DType.bool
elif isinstance(item, int):
dt = DType.int64
elif isinstance(item, float):
dt = tensorplay.get_default_dtype()
elif isinstance(item, complex):
dt = DType.complex128 if tensorplay.get_default_dtype() == DType.float64 else DType.complex64
else:
dt = _dt(item)
wrapped_result = promote(wrapped_result, dt)
def combine(higher, lower):
if higher is not None and higher.is_complex:
return higher
if lower is not None and lower.is_complex:
if higher is not None and higher.is_floating_point:
# The real tensor's width carries into the complex result.
return {
DType.float16: DType.complex32,
DType.float32: DType.complex64,
DType.float64: DType.complex128,
}.get(higher, lower)
return lower
if higher is not None and higher.is_floating_point:
return higher
if higher == DType.bool or (lower is not None and lower.is_floating_point):
return promote(higher, lower)
return higher if higher is not None else lower
return combine(dim_result, combine(zero_result, wrapped_result))
def is_nonzero(input):
n = input.numel()
if n == 0:
raise RuntimeError(
"Boolean value of Tensor with no values is ambiguous"
)
if n != 1:
raise RuntimeError(
"Boolean value of Tensor with more than one value is ambiguous"
)
return input.item() != 0
def is_same_size(input, other):
return list(input.size()) == list(other.size())
def get_device(input):
idx = getattr(input.device, "index", None)
return -1 if idx is None else int(idx)
# ---------------------------------------------------------------------------
# *_copy family
# ---------------------------------------------------------------------------
def alias_copy(input):
return input.clone()
def t_copy(input):
return input.t().clone()
def permute_copy(input, dims):
return input.permute(list(dims)).clone()
def transpose_copy(input, dim0, dim1):
return input.transpose(dim0, dim1).clone()
def squeeze_copy(input, dim=None):
if dim is None:
return input.squeeze()
return input.squeeze(dim)
def unsqueeze_copy(input, dim):
return input.unsqueeze(dim)
def select_copy(input, dim, index):
return input.select(dim, index).clone()
def slice_copy(input, dim=0, start=None, end=None, step=1):
d = _norm_dim(dim, input.dim())
n = input.size(d)
start = 0 if start is None else int(start)
end = n if end is None else int(end)
return _slice_getitem(input, d, start, end, step).contiguous()
def narrow_copy(input, dim, start, length):
return input.narrow(dim, start, length).clone()
def diagonal_copy(input, diagonal=0, dim1=-2, dim2=-1):
return input.diagonal(diagonal, dim1, dim2).clone()
def unbind_copy(input, dim=0):
return [t.clone() for t in tensorplay.unbind(input, dim)]
def split_copy(input, split_size_or_sections, dim=0):
parts = tensorplay.split(input, split_size_or_sections, dim)
return [p.clone() for p in parts]
def view_copy(input, size):
return input.view(list(size)).clone()
def reshape_as(input, other):
# input.reshape(other.shape).
return input.reshape(list(other.shape))
def unfold_copy(input, dimension, size, step):
return input.unfold(dimension, size, step).clone()
def expand_copy(input, size, *, implicit=False):
return input.expand(list(size)).clone()
def unsafe_chunk(input, chunks, dim=0):
return tensorplay.chunk(input, chunks, dim)
def unsafe_split(input, split_size, dim=0):
return tensorplay.split(input, split_size, dim)
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
def rms_norm(input, normalized_shape, weight=None, eps=None):
from tensorplay.nn import functional as F
return F.rms_norm(input, list(normalized_shape), weight, eps)
def cosine_similarity(x1, x2, dim=1, eps=1e-8):
from tensorplay.nn import functional as F
return F.cosine_similarity(x1, x2, dim, eps)
def max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1,
ceil_mode=False):
from tensorplay.nn import functional as F
return F.max_pool1d(input, kernel_size, stride=stride, padding=padding,
dilation=dilation, ceil_mode=ceil_mode)
def avg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False,
count_include_pad=True):
from tensorplay.nn import functional as F
return F.avg_pool1d(input, kernel_size, stride=stride, padding=padding,
ceil_mode=ceil_mode,
count_include_pad=count_include_pad)
def adaptive_avg_pool1d(input, output_size):
from tensorplay.nn import functional as F
return F.adaptive_avg_pool1d(input, output_size)
def adaptive_max_pool1d(input, output_size):
from tensorplay.nn import functional as F
return F.adaptive_max_pool1d(input, output_size)
# ---------------------------------------------------------------------------
# exposes at top level. Lazy import: nn.functional is not available during
# the tensorplay bootstrap (this module loads before the nn package).
# ---------------------------------------------------------------------------
def dropout(input, p=0.5, training=True, inplace=False):
from tensorplay.nn import functional as _F
return _F.dropout(input, p=p, training=training, inplace=inplace)
def dropout_(input, p=0.5, training=True):
from tensorplay.nn import functional as _F
return _F.dropout_(input, p=p, training=training)
def alpha_dropout(input, p=0.5, training=True, inplace=False):
from tensorplay.nn import functional as _F
return _F.alpha_dropout(input, p=p, training=training, inplace=inplace)
def feature_dropout(input, p=0.5, training=False, inplace=False):
from tensorplay.nn import functional as _F
return _F.feature_dropout(input, p=p, training=training,
inplace=inplace)
def feature_dropout_(input, p=0.5, training=True):
from tensorplay.nn import functional as _F
return _F.feature_dropout_(input, p=p, training=training)
def feature_alpha_dropout(input, p=0.5, training=False, inplace=False):
from tensorplay.nn import functional as _F
return _F.feature_alpha_dropout(input, p=p, training=training,
inplace=inplace)
def rrelu(input, lower=1.0 / 8, upper=1.0 / 3, training=False,
inplace=False):
from tensorplay.nn import functional as _F
return _F.rrelu(input, lower=lower, upper=upper, training=training,
inplace=inplace)
def rrelu_(input, lower=1.0 / 8, upper=1.0 / 3, training=False):
from tensorplay.nn import functional as _F
return _F.rrelu_(input, lower=lower, upper=upper, training=training)
def bilinear(input1, input2, weight, bias=None):
from tensorplay.nn import functional as _F
return _F.bilinear(input1, input2, weight, bias)
def ctc_loss(log_probs, targets, input_lengths, target_lengths,
blank=0, reduction="mean", zero_infinity=False):
from tensorplay.nn import functional as _F
return _F.ctc_loss(log_probs, targets, input_lengths, target_lengths,
blank=blank, reduction=reduction,
zero_infinity=zero_infinity)
def embedding_bag(input, weight, offsets=None, max_norm=None,
norm_type=2, scale_grad_by_freq=False, mode="mean",
sparse=False, per_sample_weights=None,
include_last_offset=False, padding_idx=None):
from tensorplay.nn import functional as _F
return _F.embedding_bag(input, weight, offsets=offsets,
max_norm=max_norm, norm_type=norm_type,
scale_grad_by_freq=scale_grad_by_freq,
mode=mode, sparse=sparse,
per_sample_weights=per_sample_weights,
include_last_offset=include_last_offset,
padding_idx=padding_idx)
def conv_tbc(input, weight, bias=None, pad=0):
from tensorplay.nn import functional as _F
return _F.conv_tbc(input, weight, bias, pad=pad)
def max_pool3d(input, kernel_size, stride=None, padding=0, dilation=1,
ceil_mode=False, return_indices=False):
from tensorplay.nn import functional as _F
return _F.max_pool3d(input, kernel_size, stride=stride, padding=padding,
dilation=dilation, ceil_mode=ceil_mode,
return_indices=return_indices)
def max_pool1d_with_indices(input, kernel_size, stride=None, padding=0,
dilation=1, ceil_mode=False):
from tensorplay.nn import functional as _F
return _F.max_pool1d_with_indices(input, kernel_size, stride=stride,
padding=padding, dilation=dilation,
ceil_mode=ceil_mode)
def native_channel_shuffle(input, groups):
from tensorplay.nn import functional as _F
return _F.channel_shuffle(input, groups)Help improve this page
Found an error, an unclear step, or a missing example?
Was this page helpful?