pub fn expand<RA: ReduceArgs, In: Numeric, Out: Numeric>(
scope: &mut Scope,
input: <RA::Input<In> as CubeType>::ExpandType,
output: <RA::Output<Out> as CubeType>::ExpandType,
) -> <(VirtualTensor<In>, VirtualTensor<Out, ReadWrite>) as CubeType>::ExpandType