Skip to content

Alloc vs BroadcastTo vs Second #367

Closed
@ricardoV94

Description

@ricardoV94

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:

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.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions