var_mean needs an axes bound check

[kevinstephano]

Currently, the var_mean API segfaults on when a dimension to reduce is asked for that is out of bounds.

import torch
from nvfuser import FusionDefinition, DataType

inputs = [
    torch.randn(256, 28, 28, 4, 32, device='cuda'),
]

def func(fd: FusionDefinition) :
    T0 = fd.from_pytorch(inputs[0])
    S1 = fd.define_scalar(1.0e-5)
    var, mean = fd.ops.var_mean(T0, axes=[1, 2, 5], correction=0, keepdim=True)
    T2 = (T0 - mean) / fd.ops.sqrt(var + S1)
    fd.add_output(T2)

with FusionDefinition() as fd:
    func(fd)

print(fd)

for _ in range(5):
    out = fd.execute(inputs)

[kevinstephano]

Closing.