mirror of
https://github.com/wassname/pytorch-ts.git
synced 2026-07-18 12:40:51 +08:00
Kashif RasulandGitHub
e685cf4b39
* ZeroInflated output added ZIP and ZINB outputs * fix import * use torch.sigmoid
31 lines
1.1 KiB
Python
31 lines
1.1 KiB
Python
# Copyright (c) 2017-2019 Uber Technologies, Inc.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
|
|
def broadcast_shape(*shapes, **kwargs):
|
|
"""
|
|
Similar to ``np.broadcast()`` but for shapes.
|
|
Equivalent to ``np.broadcast(*map(np.empty, shapes)).shape``.
|
|
|
|
:param tuple shapes: shapes of tensors.
|
|
:param bool strict: whether to use extend-but-not-resize broadcasting.
|
|
:returns: broadcasted shape
|
|
:rtype: tuple
|
|
:raises: ValueError
|
|
"""
|
|
strict = kwargs.pop("strict", False)
|
|
reversed_shape = []
|
|
for shape in shapes:
|
|
for i, size in enumerate(reversed(shape)):
|
|
if i >= len(reversed_shape):
|
|
reversed_shape.append(size)
|
|
elif reversed_shape[i] == 1 and not strict:
|
|
reversed_shape[i] = size
|
|
elif reversed_shape[i] != size and (size != 1 or strict):
|
|
raise ValueError(
|
|
"shape mismatch: objects cannot be broadcast to a single shape: {}".format(
|
|
" vs ".join(map(str, shapes))
|
|
)
|
|
)
|
|
return tuple(reversed(reversed_shape))
|