i think this is not right, and i believe it's my fault for not explaining it properly

• RFFT => mult = N/2, except for i_n = 0 and i_n = dim(i_n) - 1 whose value is then N
• IRFFT => mult = 2/N, except for i_n = 0 and i_n = dim(i_n) - 1 whose value is then 1/N

what i meant by i_n = 0 is that the n-th index to be equal to 0 and the res to be "colons", so i really meant [:,:,:,...,:,0]

but now rechecking, this is wrong because it should be for slices where i_0 = 0 and i_0 = dim(i_0) - 1. it's hard because StableHLO has weird semantics: if doing a 3-dim FFT, it performs FFT on 1st dimension and last 2 ones

so... for the case of RFFT and IRFFT, this Julia code should be equivalent:

value = if FFT
    N
elseif IFFT
    1/N
elseif RFFT
    N/2
else # IRFFT
    2/N
end

multiplier = fill(value, size(input))

if RFFT || IRFFT
    value = RFFT ? N : 1/N
    selectdim(multiplier, 1, 1) .= value
    selectdim(multiplier, 1, size(input, 1)) .= value
end

also note that because it's just a multiplier... maybe we could just skip the dynamic_update_slice in here, let it multiply with the FFT result and call slice + multiply + dynamic_update_slice on the result to correct it. this has the advantage that no array instantiation will happen even after optimizations.

hm perhaps I'll let you take it from here then. That said it is much better to have it done here (since in batched mode the constant is generated once vs in the ops itself we'd have to do it for each)