def eager_stack_homogeneous(name, *parts):
assert parts
output = parts[0].output
part_inputs = OrderedDict()
for part in parts:
assert part.output == output
assert name not in part.inputs
part_inputs.update(part.inputs)
shape = tuple(d.size for d in part_inputs.values()) + output.shape
data = ops.stack(
0, *[
ops.expand(align_tensor(part_inputs, part), shape)
for part in parts
])
inputs = OrderedDict([(name, Bint[len(parts)])])
inputs.update(part_inputs)
return Tensor(data, inputs, dtype=output.dtype)
def test_tensor_stack(n, shape, dim):
tensors = [randn(shape) for _ in range(n)]
actual = stack(tuple(Tensor(t) for t in tensors), dim=dim)
expected = Tensor(ops.stack(dim, *tensors))
assert_close(actual, expected)