-
Notifications
You must be signed in to change notification settings - Fork 137
Closed
Description
Description
Alloc
provides the same functionality as BroadcastTo
, and seems to be the default introduced in PyTensor rewrites in graphs this:
import pytensor
import pytensor.tensor as pt
x = pt.scalar("x")
out = x + [5, 5, 5]
fn = pytensor.function([x], out)
pytensor.dprint(fn, print_type=True)
Alloc [id A] <Vector(float64, shape=(3,))> 2
├─ Add [id B] <Vector(float64, shape=(1,))> 1
│ ├─ [5.] [id C] <Vector(float64, shape=(1,))>
│ └─ ExpandDims{axis=0} [id D] <Vector(float64, shape=(1,))> 0
│ └─ x [id E] <Scalar(float64, shape=())>
└─ 3 [id F] <Scalar(int64, shape=())>
This is introduced usually by this helper:
pytensor/pytensor/tensor/rewriting/basic.py
Lines 90 to 117 in 36161e8
def broadcast_like(value, template, fgraph, dtype=None): | |
""" | |
Return a Variable with the same shape and dtype as the template, | |
filled by broadcasting value through it. `value` will be cast as | |
necessary. | |
""" | |
value = as_tensor_variable(value) | |
if value.type.is_super(template.type): | |
return value | |
if template not in fgraph.variables: | |
raise NotImplementedError( | |
"broadcast_like currently requires the " | |
"template Variable to be in the fgraph already" | |
) | |
if dtype is None: | |
dtype = template.dtype | |
value = cast(value, dtype) | |
if value.type.is_super(template.type): | |
return value | |
if hasattr(fgraph, "shape_feature"): | |
new_shape = fgraph.shape_feature.shape_of[template] | |
else: | |
new_shape = template.shape | |
rval = alloc(value, *new_shape) | |
assert rval.type.dtype == dtype | |
return rval |
It doesn't make sense to have two operators for the same functionality, so we should decide which one to support.
This was added in Aesara in aesara-devs/aesara#145
The original issue mentions the alloc / vs view question: aesara-devs/aesara#36, but it seems that could easily be achieved by a single Op by manipulating the view
flag.