AI: some functions now can use lists so no need to extranally loop
This commit is contained in:
mcDandy
@@ -176,7 +176,8 @@ You can also get the node from comfy manager under the name of More math.
|
||||
|
||||
#### 1.6 Interpolation & Remap
|
||||
- `lerp(a, b, t)`: linear interpolation.
|
||||
- `slerp(a, b, t, [dim])`: spherical interpolation of vector directions with linearly interpolated magnitudes. Without `dim`, all elements form one vector. Set `dim` to select an axis of vector components; all other axes identify independent vectors, and `t` may be scalar or broadcast over those axes.
|
||||
- `slerp(a, b, t, [dim])`: spherical interpolation of vector directions with linearly interpolated magnitudes. Without `dim`, all elements form one vector. Set `dim` to select an axis of vector components; all other axes identify independent vectors, and `t` may be scalar or a tensor broadcast over those axes.
|
||||
- With a list `t`, `lerp`, `slerp` and the `*_ease` functions return one result per time (as does `interp` with a list `x`); list inputs from an earlier call are paired with their time, so nested curves work. Join the results into a batch with `concat(..., 0)`, e.g. `concat(slerp(slerp(V0, V1, w), slerp(V2, V3, w), w), 0)`.
|
||||
- `smoothstep(x, e0, e1)`: 3rd-order smooth transition.
|
||||
- `smootherstep(x, e0, e1)`: 5th-order smoother transition.
|
||||
- `cubic_ease(a, b, t)` / `cubic`: cubic easing interpolation.
|
||||
@@ -257,8 +258,8 @@ You can also get the node from comfy manager under the name of More math.
|
||||
- `append(a, b)`: append or concatenate items, lists, and tensors.
|
||||
- `repeat(x, count, [dims])`: repeat values; tensor `count` can be a scalar or a per-dimension list when `dims` is provided.
|
||||
- `batch_shuffle(tensor, indices)` / `shuffle` / `select`: reorder along batch dimension.
|
||||
- `concatenate(..., dim)` / `concat` / `cat`: concatenate tensors or lists.
|
||||
- `stack(tensors, [dim])`: stack a non-empty list of tensors along a new dimension.
|
||||
- `concatenate(..., dim)` / `concat` / `cat`: concatenate tensors or lists; a single list of tensors, `concat(list, dim)`, concatenates its elements.
|
||||
- `stack(tensors, [dim])` / `stack(x, ..., dim)`: stack a non-empty list of tensors, or separate tensors followed by `dim`, along a new dimension.
|
||||
- `roll(tensor, shifts, [dim])`: circular shift.
|
||||
- `tensor(shape, [value, [type]])`: create a filled tensor; `type` a tensor to copy its dtype to self when being created.
|
||||
- `zeros(shape, [dtype_template])` / `ones(shape, [dtype_template])`: create zero- or one-filled tensors; an optional tensor supplies the output dtype.
|
||||
|
||||
@@ -1574,16 +1574,17 @@ funcN:
|
||||
| UNSQUEEZE LPAREN expr COMMA expr RPAREN # UnsqueezeFunc
|
||||
/**
|
||||
concatenate(x, ..., dim) - concatenates one or more tensors along the specified dimension
|
||||
- x: first tensor or list; more may follow before dim
|
||||
- x: first tensor or list; more may follow before dim. A single list of tensors joins its elements
|
||||
- dim: dimension to join along (the last argument)
|
||||
*/
|
||||
| CONCAT LPAREN expr (COMMA expr)+ RPAREN # ConcatFunc
|
||||
/**
|
||||
stack(tensors, [dim]) - stacks a list of tensors along a new dimension
|
||||
stack(tensors, [dim]) / stack(x, ..., dim) - stacks tensors along a new dimension
|
||||
- tensors: list of tensors with equal shapes
|
||||
- dim: position of the new dimension (default: 0)
|
||||
- x: first tensor; more may follow before dim
|
||||
- dim: position of the new dimension (default: 0 for a list; required after separate tensors)
|
||||
*/
|
||||
| STACK LPAREN expr (COMMA expr)? RPAREN # StackFunc
|
||||
| STACK LPAREN expr (COMMA expr)* RPAREN # StackFunc
|
||||
/**
|
||||
nonzero(x) - returns an N-by-rank tensor of coordinates for nonzero elements
|
||||
- x: tensor
|
||||
|
||||
@@ -1737,19 +1737,10 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
xs = [k for k, _ in keys]
|
||||
ys = [v for _, v in keys]
|
||||
|
||||
if self._is_list(x):
|
||||
return [self._interp_at(float(v), xs, ys, mode) for v in x]
|
||||
if not self._is_tensor(x) or x.numel() == 1:
|
||||
i = bisect.bisect_right(xs, float(x))
|
||||
if i == 0:
|
||||
return ys[0]
|
||||
if i == len(xs) or mode == "step":
|
||||
return ys[i - 1]
|
||||
t = (float(x) - xs[i - 1]) / (xs[i] - xs[i - 1])
|
||||
if mode == "smooth":
|
||||
t = t * t * (3 - 2 * t)
|
||||
a, b = ys[i - 1], ys[i]
|
||||
if self._is_tensor(a) or self._is_tensor(b):
|
||||
return torch.lerp(self._promote_to_tensor(a), self._promote_to_tensor(b), t)
|
||||
return a + (b - a) * t
|
||||
return self._interp_at(float(x), xs, ys, mode)
|
||||
|
||||
if any(self._is_tensor(v) and v.numel() != 1 for v in ys):
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: interp values must be numbers when x is a tensor")
|
||||
@@ -1768,6 +1759,20 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
t = t * t * (3 - 2 * t)
|
||||
return torch.lerp(ys_t[i - 1], ys_t[i], t)
|
||||
|
||||
def _interp_at(self, x, xs, ys, mode):
|
||||
i = bisect.bisect_right(xs, x)
|
||||
if i == 0:
|
||||
return ys[0]
|
||||
if i == len(xs) or mode == "step":
|
||||
return ys[i - 1]
|
||||
t = (x - xs[i - 1]) / (xs[i] - xs[i - 1])
|
||||
if mode == "smooth":
|
||||
t = t * t * (3 - 2 * t)
|
||||
a, b = ys[i - 1], ys[i]
|
||||
if self._is_tensor(a) or self._is_tensor(b):
|
||||
return torch.lerp(self._promote_to_tensor(a), self._promote_to_tensor(b), t)
|
||||
return a + (b - a) * t
|
||||
|
||||
def _schedule_args(self, ctx):
|
||||
steps = self._to_int((yield ctx.expr(0)), ctx, "schedule steps", strict=True)
|
||||
return steps, float((yield ctx.expr(1))), float((yield ctx.expr(2)))
|
||||
@@ -1799,17 +1804,35 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
|
||||
if self._is_math_dict(a) or self._is_math_dict(b) or self._is_math_dict(w):
|
||||
return self._map_math_dict(self._lerp_helper, a, b, w)
|
||||
return self._lerp_times(a, b, w)
|
||||
|
||||
def _lerp_times(self, a, b, w):
|
||||
if self._is_list(w):
|
||||
return [self._lerp_helper(a[i] if self._is_list(a) else a,
|
||||
b[i] if self._is_list(b) else b,
|
||||
t) for i, t in enumerate(w)]
|
||||
|
||||
return self._lerp_helper(a, b, w)
|
||||
|
||||
def visitSlerpFunc(self, ctx):
|
||||
a = self._promote_to_tensor((yield ctx.expr(0)))
|
||||
b = self._promote_to_tensor((yield ctx.expr(1)))
|
||||
t = self._promote_to_tensor((yield ctx.expr(2)))
|
||||
a = (yield ctx.expr(0))
|
||||
b = (yield ctx.expr(1))
|
||||
t = (yield ctx.expr(2))
|
||||
dim = (yield ctx.expr(3)) if len(ctx.expr()) > 3 else None
|
||||
if not self._is_list(t):
|
||||
return self._slerp_helper(ctx, a, b, t, dim)
|
||||
# a list t evaluates every time in it; lists of tensors (from an earlier list slerp/lerp) are paired with it
|
||||
def at(x, i):
|
||||
if self._is_list(x) and any(self._is_tensor(v) for v in x):
|
||||
if len(x) != len(t):
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: slerp got {len(x)} values for {len(t)} times")
|
||||
return x[i]
|
||||
return x
|
||||
return [self._slerp_helper(ctx, at(a, i), at(b, i), ti, dim) for i, ti in enumerate(t)]
|
||||
|
||||
def _slerp_helper(self, ctx, a, b, t, dim):
|
||||
a = self._promote_to_tensor(a)
|
||||
b = self._promote_to_tensor(b)
|
||||
t = self._promote_to_tensor(t)
|
||||
if a.ndim == 0 or b.ndim == 0:
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: slerp inputs must have a vector dimension")
|
||||
try:
|
||||
@@ -1817,14 +1840,14 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
except RuntimeError as exc:
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: slerp inputs must be broadcastable") from exc
|
||||
|
||||
default_dim = len(ctx.expr()) <= 3
|
||||
default_dim = dim is None
|
||||
if default_dim:
|
||||
output_shape = a.shape
|
||||
a = a.reshape(-1)
|
||||
b = b.reshape(-1)
|
||||
dim = 0
|
||||
else:
|
||||
dims = self._normalize_dims((yield ctx.expr(3)), a.ndim, ctx, "slerp dimension")
|
||||
dims = self._normalize_dims(dim, a.ndim, ctx, "slerp dimension")
|
||||
if len(dims) != 1:
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: slerp expects one vector dimension")
|
||||
dim = dims[0]
|
||||
@@ -1992,13 +2015,13 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
def cubic(v):
|
||||
return torch.where(v < 0.5, 4 * v**3, 1 - torch.pow(-2 * v + 2, 3) / 2) if self._is_tensor(v) else \
|
||||
(4 * v**3 if v < 0.5 else 1 - math.pow(-2 * v + 2, 3) / 2)
|
||||
return self._lerp_helper(a, b, cubic(t))
|
||||
return self._lerp_times(a, b, [cubic(v) for v in t] if self._is_list(t) else cubic(t))
|
||||
|
||||
def visitSineEaseFunc(self, ctx):
|
||||
a, b, t = (yield ctx.expr(0)), (yield ctx.expr(1)), (yield ctx.expr(2))
|
||||
def sine(v):
|
||||
return -(torch.cos(math.pi * v) - 1) / 2 if self._is_tensor(v) else -(math.cos(math.pi * v) - 1) / 2
|
||||
return self._lerp_helper(a, b, sine(t))
|
||||
return self._lerp_times(a, b, [sine(v) for v in t] if self._is_list(t) else sine(t))
|
||||
|
||||
def visitElasticEaseFunc(self, ctx):
|
||||
a, b, t = (yield ctx.expr(0)), (yield ctx.expr(1)), (yield ctx.expr(2))
|
||||
@@ -2013,7 +2036,7 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
if v >= 1: return 1
|
||||
if v < 0.5: return -(math.pow(2, 20 * v - 10) * math.sin((20 * v - 11.125) * c4)) / 2
|
||||
return (math.pow(2, -20 * v + 10) * math.sin((20 * v - 11.125) * c4)) / 2 + 1
|
||||
return self._lerp_helper(a, b, elastic(t))
|
||||
return self._lerp_times(a, b, [elastic(v) for v in t] if self._is_list(t) else elastic(t))
|
||||
|
||||
# Helpers for visiting generic exprs
|
||||
def visitFunc1Exp(self, ctx):
|
||||
@@ -5471,6 +5494,9 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
items.append((yield exprs[i]))
|
||||
|
||||
dim_val = yield exprs[-1]
|
||||
# concat(list, dim) joins the list's elements, e.g. the per-time results of lerp/slerp
|
||||
if len(items) == 1 and self._is_list(items[0]) and any(self._is_tensor(x) for x in items[0]):
|
||||
items = list(items[0])
|
||||
if any(isinstance(x, str) for x in items):
|
||||
return "".join(str(x) for x in items)
|
||||
|
||||
@@ -5485,12 +5511,22 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
return torch.cat(tensors, dim=d)
|
||||
|
||||
def visitStackFunc(self, ctx):
|
||||
values = yield ctx.expr(0)
|
||||
exprs = ctx.expr()
|
||||
values = yield exprs[0]
|
||||
dim = 0
|
||||
if self._is_list(values) and len(exprs) <= 2:
|
||||
if len(exprs) > 1:
|
||||
dim = self._to_int((yield exprs[1]), ctx, "stack dim", strict=True)
|
||||
elif len(exprs) > 1:
|
||||
# stack(x, ..., dim), like concat
|
||||
values = [values]
|
||||
for e in exprs[1:-1]:
|
||||
values.append((yield e))
|
||||
dim = self._to_int((yield exprs[-1]), ctx, "stack dim", strict=True)
|
||||
if not self._is_list(values) or not values:
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: stack expects a non-empty list of values")
|
||||
raise ValueError(f"{ctx.start.line}:{ctx.start.column}: stack expects a non-empty list of values, or values followed by dim")
|
||||
|
||||
tensors = [self._promote_to_tensor(value) for value in values]
|
||||
dim = self._to_int((yield ctx.expr(1)), ctx, "stack dim", strict=True) if len(ctx.expr()) > 1 else 0
|
||||
if dim < 0:
|
||||
dim += tensors[0].ndim + 1
|
||||
if dim < 0 or dim > tensors[0].ndim:
|
||||
|
||||
File diff suppressed because one or more lines are too long
+1073
-1066
File diff suppressed because it is too large
Load Diff
@@ -228,10 +228,10 @@ INBUILT_FUNCTION_META = {
|
||||
'stack_clear': {'min_args': 1, 'max_args': 1, 'signature': 'stack_clear(slot)', 'forms': [['slot']], 'description': 'clears the contents of slot stack', 'details': '', 'params': {'slot': 'stack number (any integer)'}, 'aliases': []},
|
||||
'stack_has': {'min_args': 1, 'max_args': 1, 'signature': 'stack_has(slot)', 'forms': [['slot']], 'description': 'returns 1.0 if slot exists in stack and is not empty else returns 0.0', 'details': '', 'params': {'slot': 'stack number (any integer)'}, 'aliases': []},
|
||||
'stack_get': {'min_args': 1, 'max_args': 1, 'signature': 'stack_get(slot)', 'forms': [['slot']], 'description': 'returns the last pushed value in slot without popping it', 'details': '', 'params': {'slot': 'stack number (any integer)'}, 'aliases': []},
|
||||
'cat': {'min_args': 2, 'max_args': None, 'signature': 'cat(x, ..., dim)', 'forms': [['x', '...', 'dim']], 'description': 'concatenates one or more tensors along the specified dimension', 'details': '', 'params': {'x': 'first tensor or list; more may follow before dim', 'dim': 'dimension to join along (the last argument)'}, 'aliases': ['concat', 'concatenate']},
|
||||
'concat': {'min_args': 2, 'max_args': None, 'signature': 'concat(x, ..., dim)', 'forms': [['x', '...', 'dim']], 'description': 'concatenates one or more tensors along the specified dimension', 'details': '', 'params': {'x': 'first tensor or list; more may follow before dim', 'dim': 'dimension to join along (the last argument)'}, 'aliases': ['cat', 'concatenate']},
|
||||
'concatenate': {'min_args': 2, 'max_args': None, 'signature': 'concatenate(x, ..., dim)', 'forms': [['x', '...', 'dim']], 'description': 'concatenates one or more tensors along the specified dimension', 'details': '', 'params': {'x': 'first tensor or list; more may follow before dim', 'dim': 'dimension to join along (the last argument)'}, 'aliases': ['cat', 'concat']},
|
||||
'stack': {'min_args': 1, 'max_args': 2, 'signature': 'stack(tensors, [dim])', 'forms': [['tensors', '[dim]']], 'description': 'stacks a list of tensors along a new dimension', 'details': '', 'params': {'tensors': 'list of tensors with equal shapes', 'dim': 'position of the new dimension (default: 0)'}, 'aliases': []},
|
||||
'cat': {'min_args': 2, 'max_args': None, 'signature': 'cat(x, ..., dim)', 'forms': [['x', '...', 'dim']], 'description': 'concatenates one or more tensors along the specified dimension', 'details': '', 'params': {'x': 'first tensor or list; more may follow before dim. A single list of tensors joins its elements', 'dim': 'dimension to join along (the last argument)'}, 'aliases': ['concat', 'concatenate']},
|
||||
'concat': {'min_args': 2, 'max_args': None, 'signature': 'concat(x, ..., dim)', 'forms': [['x', '...', 'dim']], 'description': 'concatenates one or more tensors along the specified dimension', 'details': '', 'params': {'x': 'first tensor or list; more may follow before dim. A single list of tensors joins its elements', 'dim': 'dimension to join along (the last argument)'}, 'aliases': ['cat', 'concatenate']},
|
||||
'concatenate': {'min_args': 2, 'max_args': None, 'signature': 'concatenate(x, ..., dim)', 'forms': [['x', '...', 'dim']], 'description': 'concatenates one or more tensors along the specified dimension', 'details': '', 'params': {'x': 'first tensor or list; more may follow before dim. A single list of tensors joins its elements', 'dim': 'dimension to join along (the last argument)'}, 'aliases': ['cat', 'concat']},
|
||||
'stack': {'min_args': 1, 'max_args': None, 'signature': 'stack(tensors, [dim]) / stack(x, ..., dim)', 'forms': [['tensors', '[dim]'], ['x', '...', 'dim']], 'description': 'stacks tensors along a new dimension', 'details': '', 'params': {'tensors': 'list of tensors with equal shapes', 'x': 'first tensor; more may follow before dim', 'dim': 'position of the new dimension (default: 0 for a list; required after separate tensors)'}, 'aliases': []},
|
||||
'nonzero': {'min_args': 1, 'max_args': 1, 'signature': 'nonzero(x)', 'forms': [['x']], 'description': 'returns an N-by-rank tensor of coordinates for nonzero elements', 'details': '', 'params': {'x': 'tensor'}, 'aliases': []},
|
||||
'int': {'min_args': 1, 'max_args': 1, 'signature': 'int(x)', 'forms': [['x']], 'description': 'casts x to integer', 'details': '', 'params': {'x': 'value to convert; lists convert each item and text is parsed'}, 'aliases': []},
|
||||
'float': {'min_args': 1, 'max_args': 1, 'signature': 'float(x)', 'forms': [['x']], 'description': 'casts x to float', 'details': '', 'params': {'x': 'value to convert; lists convert each item and text is parsed'}, 'aliases': []},
|
||||
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
@@ -240,7 +240,7 @@ def test_slerp_validates_broadcasting_and_preserves_float8_dtype():
|
||||
with pytest.raises(ValueError, match="inputs must be broadcastable"):
|
||||
eval_expr("slerp([1, 2], [1, 2, 3], 0.5)")
|
||||
with pytest.raises(ValueError, match="t must broadcast"):
|
||||
eval_expr("slerp([[1, 0], [0, 1]], [[0, 1], [1, 0]], [1, 2, 3])")
|
||||
eval_expr("slerp([[1, 0], [0, 1]], [[0, 1], [1, 0]], t, 1)", {"t": torch.tensor([1.0, 2.0, 3.0])})
|
||||
|
||||
float8_dtypes = [
|
||||
getattr(torch, name, None)
|
||||
@@ -251,6 +251,56 @@ def test_slerp_validates_broadcasting_and_preserves_float8_dtype():
|
||||
result = eval_expr("slerp(a, b, 0.5)", {"a": values, "b": values})
|
||||
assert result.dtype == dtype
|
||||
|
||||
def test_slerp_and_lerp_with_list_of_times_concat_to_batch():
|
||||
a = torch.randn(1, 3, 4)
|
||||
b = torch.randn(1, 3, 4)
|
||||
c = torch.randn(1, 3, 4)
|
||||
times = [0.0, 0.25, 1.0]
|
||||
variables = {"a": a, "b": b, "c": c}
|
||||
|
||||
result = eval_expr("concat(slerp(a, b, [0, 0.25, 1]), 0)", variables)
|
||||
expected = torch.cat([eval_expr(f"slerp(a, b, {t})", variables) for t in times])
|
||||
torch.testing.assert_close(result, expected)
|
||||
torch.testing.assert_close(result[0:1], a)
|
||||
torch.testing.assert_close(result[2:3], b)
|
||||
|
||||
# nested calls pair each time with its own result
|
||||
result = eval_expr("w = [0, 0.25, 1]; concat(slerp(slerp(a, b, w), slerp(b, c, w), w), 0)", variables)
|
||||
expected = torch.cat([eval_expr(f"slerp(slerp(a, b, {t}), slerp(b, c, {t}), {t})", variables) for t in times])
|
||||
torch.testing.assert_close(result, expected)
|
||||
|
||||
result = eval_expr("w = [0, 0.25, 1]; concat(lerp(lerp(a, b, w), lerp(b, c, w), w), 0)", variables)
|
||||
expected = torch.cat([torch.lerp(torch.lerp(a, b, t), torch.lerp(b, c, t), t) for t in times])
|
||||
torch.testing.assert_close(result, expected)
|
||||
|
||||
with pytest.raises(ValueError, match="2 values for 3 times"):
|
||||
eval_expr("slerp(slerp(a, b, [0, 1]), b, [0, 0.5, 1])", variables)
|
||||
|
||||
def test_eases_and_interp_with_list_of_times():
|
||||
a = torch.randn(1, 3, 4)
|
||||
b = torch.randn(1, 3, 4)
|
||||
variables = {"a": a, "b": b}
|
||||
times = [0.0, 0.3, 0.8, 1.0]
|
||||
for name in ("cubic_ease", "sine_ease", "elastic_ease"):
|
||||
result = eval_expr(f"concat({name}(a, b, [0, 0.3, 0.8, 1]), 0)", variables)
|
||||
expected = torch.cat([eval_expr(f"{name}(a, b, {t})", variables) for t in times])
|
||||
torch.testing.assert_close(result, expected)
|
||||
assert eval_expr("sine_ease(0, 2, [0, 0.5, 1])") == pytest.approx([0.0, 1.0, 2.0])
|
||||
|
||||
keys = {"a": a, "b": b, "c": torch.randn(1, 3, 4)}
|
||||
for mode in ("linear", "smooth", "step"):
|
||||
result = eval_expr(f'concat(interp([0, 0.25, 0.5, 0.9], [0, 0.5, 1], [a, b, c], "{mode}"), 0)', keys)
|
||||
expected = torch.cat([eval_expr(f'interp({t}, [0, 0.5, 1], [a, b, c], "{mode}")', keys) for t in (0, 0.25, 0.5, 0.9)])
|
||||
torch.testing.assert_close(result, expected)
|
||||
assert eval_expr("interp([0, 12, 24], [0, 24], [1, 3])") == [1.0, 2.0, 3.0]
|
||||
|
||||
def test_concat_single_list_argument():
|
||||
a = torch.ones(1, 2)
|
||||
b = torch.zeros(2, 2)
|
||||
torch.testing.assert_close(eval_expr("concat([a, b], 0)", {"a": a, "b": b}), torch.cat([a, b]))
|
||||
torch.testing.assert_close(eval_expr("concat([a, a], 1)", {"a": a}), torch.ones(1, 4))
|
||||
assert eval_expr("concat([1, 2], [3], 0)") == [1, 2, 3]
|
||||
|
||||
def test_sinc_and_erfc():
|
||||
assert {"sinc", "erfc"} <= INBUILT_FUNCTIONS
|
||||
assert eval_expr("sinc(0)") == 1.0
|
||||
|
||||
@@ -111,6 +111,17 @@ def test_stack_list_of_tensors_with_negative_dimension():
|
||||
assert torch.equal(result, torch.tensor([[1.0, 3.0], [2.0, 4.0]]))
|
||||
|
||||
|
||||
def test_stack_separate_tensors_followed_by_dimension():
|
||||
first = torch.tensor([1.0, 2.0])
|
||||
second = torch.tensor([3.0, 4.0])
|
||||
variables = {"a": first, "b": second}
|
||||
|
||||
assert torch.equal(_visit("stack(a, b, 0)", variables), torch.stack([first, second]))
|
||||
assert torch.equal(_visit("stack(a, b, a, -1)", variables), torch.stack([first, second, first], dim=-1))
|
||||
assert torch.equal(_visit("stack(a, 0)", variables), first.unsqueeze(0))
|
||||
assert torch.equal(_visit("stack([a, b], 1)", variables), torch.stack([first, second], dim=1))
|
||||
|
||||
|
||||
def test_stack_rejects_invalid_dimension_and_non_list_input():
|
||||
with pytest.raises(ValueError, match="stack expects a non-empty list"):
|
||||
_visit("stack(a)", {"a": torch.tensor([1.0, 2.0])})
|
||||
|
||||
@@ -231,10 +231,10 @@ export const FUNCTION_META = {
|
||||
"stack_clear": { minArgs: 1, maxArgs: 1, signature: "stack_clear(slot)", forms: [["slot"]], description: "clears the contents of slot stack", details: "", params: {"slot": "stack number (any integer)"}, aliases: [] },
|
||||
"stack_has": { minArgs: 1, maxArgs: 1, signature: "stack_has(slot)", forms: [["slot"]], description: "returns 1.0 if slot exists in stack and is not empty else returns 0.0", details: "", params: {"slot": "stack number (any integer)"}, aliases: [] },
|
||||
"stack_get": { minArgs: 1, maxArgs: 1, signature: "stack_get(slot)", forms: [["slot"]], description: "returns the last pushed value in slot without popping it", details: "", params: {"slot": "stack number (any integer)"}, aliases: [] },
|
||||
"cat": { minArgs: 2, maxArgs: null, signature: "cat(x, ..., dim)", forms: [["x", "...", "dim"]], description: "concatenates one or more tensors along the specified dimension", details: "", params: {"x": "first tensor or list; more may follow before dim", "dim": "dimension to join along (the last argument)"}, aliases: ["concat", "concatenate"] },
|
||||
"concat": { minArgs: 2, maxArgs: null, signature: "concat(x, ..., dim)", forms: [["x", "...", "dim"]], description: "concatenates one or more tensors along the specified dimension", details: "", params: {"x": "first tensor or list; more may follow before dim", "dim": "dimension to join along (the last argument)"}, aliases: ["cat", "concatenate"] },
|
||||
"concatenate": { minArgs: 2, maxArgs: null, signature: "concatenate(x, ..., dim)", forms: [["x", "...", "dim"]], description: "concatenates one or more tensors along the specified dimension", details: "", params: {"x": "first tensor or list; more may follow before dim", "dim": "dimension to join along (the last argument)"}, aliases: ["cat", "concat"] },
|
||||
"stack": { minArgs: 1, maxArgs: 2, signature: "stack(tensors, [dim])", forms: [["tensors", "[dim]"]], description: "stacks a list of tensors along a new dimension", details: "", params: {"tensors": "list of tensors with equal shapes", "dim": "position of the new dimension (default: 0)"}, aliases: [] },
|
||||
"cat": { minArgs: 2, maxArgs: null, signature: "cat(x, ..., dim)", forms: [["x", "...", "dim"]], description: "concatenates one or more tensors along the specified dimension", details: "", params: {"x": "first tensor or list; more may follow before dim. A single list of tensors joins its elements", "dim": "dimension to join along (the last argument)"}, aliases: ["concat", "concatenate"] },
|
||||
"concat": { minArgs: 2, maxArgs: null, signature: "concat(x, ..., dim)", forms: [["x", "...", "dim"]], description: "concatenates one or more tensors along the specified dimension", details: "", params: {"x": "first tensor or list; more may follow before dim. A single list of tensors joins its elements", "dim": "dimension to join along (the last argument)"}, aliases: ["cat", "concatenate"] },
|
||||
"concatenate": { minArgs: 2, maxArgs: null, signature: "concatenate(x, ..., dim)", forms: [["x", "...", "dim"]], description: "concatenates one or more tensors along the specified dimension", details: "", params: {"x": "first tensor or list; more may follow before dim. A single list of tensors joins its elements", "dim": "dimension to join along (the last argument)"}, aliases: ["cat", "concat"] },
|
||||
"stack": { minArgs: 1, maxArgs: null, signature: "stack(tensors, [dim]) / stack(x, ..., dim)", forms: [["tensors", "[dim]"], ["x", "...", "dim"]], description: "stacks tensors along a new dimension", details: "", params: {"tensors": "list of tensors with equal shapes", "x": "first tensor; more may follow before dim", "dim": "position of the new dimension (default: 0 for a list; required after separate tensors)"}, aliases: [] },
|
||||
"nonzero": { minArgs: 1, maxArgs: 1, signature: "nonzero(x)", forms: [["x"]], description: "returns an N-by-rank tensor of coordinates for nonzero elements", details: "", params: {"x": "tensor"}, aliases: [] },
|
||||
"int": { minArgs: 1, maxArgs: 1, signature: "int(x)", forms: [["x"]], description: "casts x to integer", details: "", params: {"x": "value to convert; lists convert each item and text is parsed"}, aliases: [] },
|
||||
"float": { minArgs: 1, maxArgs: 1, signature: "float(x)", forms: [["x"]], description: "casts x to float", details: "", params: {"x": "value to convert; lists convert each item and text is parsed"}, aliases: [] },
|
||||
|
||||
Reference in New Issue
Block a user