//! Demonstrates broadcasting tensors to different sizes, and axis reductions
//! with BroadcastTo and ReduceTo

use dfdx::{
    shapes::{Axis, Rank2, Rank4},
    tensor::{AsArray, AutoDevice, TensorFrom},
    tensor_ops::{BroadcastTo, MeanTo},
};

fn main() {
    let dev = AutoDevice::default();
    let a = dev.tensor([1.0f32, 2.0, 3.0]);
    // NOTE: Cuda currently does not support broadcasting.
    // Its usage results in errors and wrong outputs.

    // to broadcast, use `Broadcast::broadcast()` and specify
    // the output type. the axes that are broadcast are inferred for you!
    let b = a.broadcast::<Rank2<5, 3>, _>();
    assert_eq!(b.array(), [[1.0, 2.0, 3.0]; 5]);

    // we can really broadcast any axes on either side
    // here a (5,3) tensor is broadcast to (7,5,3,2).
    // so 7 is added in front, and 2 is added last
    let c = b.broadcast::<Rank4<7, 5, 3, 2>, _>();
    assert_eq!(c.array(), [[[[1.0; 2], [2.0; 2], [3.0; 2]]; 5]; 7]);

    // the opposite of broadcast is reducing
    // we've already introduced one reduction which is mean
    let d = c.mean::<Rank2<5, 3>, _>();
    assert_eq!(d.array(), [[1.0, 2.0, 3.0]; 5]);

    // Sometimes it's ambiguous which axes you mean to broadcast or reduce.
    // Here rust doesn't know if the new axis is the first or second.
    // We can specify the axes parameter of broadcast in addition to the new shape
    let e = dev.tensor([1.0]);
    let f = e.broadcast::<Rank2<1, 1>, Axis<1>>();
    // NOTE: will fail with "Multiple impls satisfying...":
    // let f = e.broadcast::<Rank2<1, 1>, _>();

    // reductions have the same problem when it's ambiguous,
    // so we can also use mean_along with an axis
    let _ = f.mean::<_, Axis<0>>();
}