Adding basic functionality to specify the feature axis of the input tensor (setting the default to -1)