Infer `RandomVariable` broadcastable pattern when `size` has a `Cast`

[brandonwillard]

In the following example, the broadcast pattern isn't inferred because of a Cast in size:

import aesara.tensor as at
import aesara.tensor.random.basic as ar


size_at = at.cast([3.0, 1.0], "int64")
# or
# size_at = at.cast([at.scalar(), 1.0], "int64")

norm_at = ar.normal(0, 1, size=size_at)

>>> norm_at.type
TensorType(float64, matrix)

One solution could involve lifting the Cast Op, or—for this exact case—canonicalizing the graph for size would remove the Cast altogether.

[twiecki]

@brandonwillard Can you provide a bit more detail on what needs to happen here?

[brandonwillard]

@brandonwillard Can you provide a bit more detail on what needs to happen here?

This could be accomplished in more than a few ways, but perhaps the simplest involves adding Cast to the isinstance here.

[michaelosthege]

After digging into RandomVariable.compute_bcast as @brandonwillard suggested, I don't think the if clause is even necessary:

if s_owner and isinstance(s_owner.op, Assert):
    s = s_owner.inputs[0]

try:
    s_val = get_scalar_constant_value(s)
except NotScalarConstantError:
    s_val = False

Note that get_scalar_constant_value already implements quite a bit of recursion that, among other things, also looks for Assert..

The _scalar_constant_value_elemwise_ops also include aes.Cast, but for some reason this part doesn't work as expected:

const = at.as_tensor_variable([at.scalar(dtype="int32"), np.int32(1)])
aesara.printing.debugprint(const)
at.get_scalar_constant_value(const[1])

MakeVector{dtype='int32'} [id A] ''   
 |<TensorType(int32, scalar)> [id B]
 |TensorConstant{1} [id C]
array(1)

↑ Constant was successfully extracted 🎉
↓ But wrapping in a simple at.cast causes problems 🐛

casted = at.cast(const, dtype="int64")
aesara.printing.debugprint(casted)
at.get_scalar_constant_value(casted[1])

Elemwise{Cast{int64}} [id A] ''   
 |MakeVector{dtype='int32'} [id B] ''   
   |<TensorType(int32, scalar)> [id C]
   |TensorConstant{1} [id D]
---------------------------------------------------------------------------
NotScalarConstantError                    Traceback (most recent call last)
<ipython-input-87-a4eae33a03f7> in <module>
      1 casted = at.cast(const, dtype="int64")
      2 aesara.printing.debugprint(casted)
----> 3 at.get_scalar_constant_value(casted[1])

e:\source\repos\aesara\aesara\tensor\basic.py in get_scalar_constant_value(orig_v, elemwise, only_process_constants, max_recur)
    566                         return np.asarray(np.shape(grandparent.data)[idx])
    567 
--> 568         raise NotScalarConstantError(v)
    569 
    570 

NotScalarConstantError: Subtensor{int64}.0